QodeAssist/sources/ChatView/ChatModel.cpp

// Copyright (C) 2024-2026 Petr Mironychev
// SPDX-License-Identifier: GPL-3.0-or-later
// Additional attribution terms under GPLv3 §7(b) apply — see LICENSE

#include "ChatModel.hpp"

#include <QDir>
#include <QFileInfo>
#include <QRegularExpression>
#include <QUrl>

#include <algorithm>
#include <tuple>

#include "logger/Logger.hpp"

namespace QodeAssist::Chat {

namespace {

auto usageOf(const ChatModel::Message &message)
{
    return std::tie(
        message.promptTokens,
        message.completionTokens,
        message.cachedPromptTokens,
        message.reasoningTokens);
}

bool carriesUsage(const ChatModel::Message &message)
{
    return message.promptTokens != 0 || message.completionTokens != 0
           || message.cachedPromptTokens != 0 || message.reasoningTokens != 0;
}

} // namespace

ChatModel::ChatModel(QObject *parent)
    : QAbstractListModel(parent)
{}

int ChatModel::rowCount(const QModelIndex &parent) const
{
    return m_messages.size();
}

QVariant ChatModel::data(const QModelIndex &index, int role) const
{
    if (!index.isValid() || index.row() >= m_messages.size())
        return QVariant();

    const Message &message = m_messages[index.row()];
    switch (static_cast<Roles>(role)) {
    case Roles::RoleType:
        return QVariant::fromValue(message.role);
    case Roles::Content: {
        return message.content;
    }
    case Roles::Attachments: {
        QVariantList attachmentsList;
        for (const auto &attachment : message.attachments) {
            QVariantMap attachmentMap;
            attachmentMap["fileName"] = attachment.filename;
            attachmentMap["storedPath"] = attachment.content;
            
            if (!m_chatFilePath.isEmpty()) {
                QFileInfo fileInfo(m_chatFilePath);
                QString baseName = fileInfo.completeBaseName();
                QString dirPath = fileInfo.absolutePath();
                QString contentFolder = QDir(dirPath).filePath(baseName + "_content");
                QString fullPath = QDir(contentFolder).filePath(attachment.content);
                attachmentMap["filePath"] = fullPath;
            } else {
                attachmentMap["filePath"] = QString();
            }
            
            attachmentsList.append(attachmentMap);
        }
        return attachmentsList;
    }
    case Roles::IsRedacted: {
        return message.isRedacted;
    }
    case Roles::PromptTokens:
        return message.promptTokens;
    case Roles::CompletionTokens:
        return message.completionTokens;
    case Roles::CachedPromptTokens:
        return message.cachedPromptTokens;
    case Roles::ReasoningTokens:
        return message.reasoningTokens;
    case Roles::TotalTokens:
        return message.promptTokens + message.completionTokens;
    case Roles::Images: {
        QVariantList imagesList;
        for (const auto &image : message.images) {
            QVariantMap imageMap;
            imageMap["fileName"] = image.fileName;
            imageMap["storedPath"] = image.storedPath;
            imageMap["mediaType"] = image.mediaType;
            
            if (!m_chatFilePath.isEmpty()) {
                QFileInfo fileInfo(m_chatFilePath);
                QString baseName = fileInfo.completeBaseName();
                QString dirPath = fileInfo.absolutePath();
                QString contentFolder = QDir(dirPath).filePath(baseName + "_content");
                QString fullPath = QDir(contentFolder).filePath(image.storedPath);
                imageMap["imageUrl"] = QUrl::fromLocalFile(fullPath).toString();
                imageMap["filePath"] = fullPath;
            } else {
                imageMap["imageUrl"] = QString();
                imageMap["filePath"] = QString();
            }
            
            imagesList.append(imageMap);
        }
        return imagesList;
    }
    default:
        return QVariant();
    }
}

QHash<int, QByteArray> ChatModel::roleNames() const
{
    QHash<int, QByteArray> roles;
    roles[Roles::RoleType] = "roleType";
    roles[Roles::Content] = "content";
    roles[Roles::Attachments] = "attachments";
    roles[Roles::IsRedacted] = "isRedacted";
    roles[Roles::Images] = "images";
    roles[Roles::PromptTokens] = "promptTokens";
    roles[Roles::CompletionTokens] = "completionTokens";
    roles[Roles::CachedPromptTokens] = "cachedPromptTokens";
    roles[Roles::ReasoningTokens] = "reasoningTokens";
    roles[Roles::TotalTokens] = "totalTokens";
    return roles;
}

void ChatModel::resetMessages(const QVector<Message> &messages)
{
    beginResetModel();
    m_messages = messages;
    endResetModel();
    emit modelReseted();
    emit sessionUsageChanged();
}

void ChatModel::appendMessages(const QVector<Message> &messages)
{
    if (messages.isEmpty())
        return;

    beginInsertRows(QModelIndex(), m_messages.size(), m_messages.size() + messages.size() - 1);
    m_messages.append(messages);
    endInsertRows();

    if (std::any_of(messages.cbegin(), messages.cend(), carriesUsage))
        emit sessionUsageChanged();
}

void ChatModel::updateMessage(int index, const Message &message)
{
    if (index < 0 || index >= m_messages.size()) {
        LOG_MESSAGE(QString("Session/model desync: update of row %1 with %2 rows present")
                        .arg(index)
                        .arg(m_messages.size()));
        return;
    }

    const bool usageChanged = usageOf(m_messages[index]) != usageOf(message);
    m_messages[index] = message;
    emit dataChanged(this->index(index), this->index(index));

    if (usageChanged)
        emit sessionUsageChanged();
}

void ChatModel::removeMessages(int first, int count)
{
    if (first < 0 || count <= 0 || first + count > m_messages.size()) {
        LOG_MESSAGE(QString("Session/model desync: removal of %1 rows at %2 with %3 rows present")
                        .arg(count)
                        .arg(first)
                        .arg(m_messages.size()));
        return;
    }

    const bool usageChanged
        = std::any_of(m_messages.cbegin() + first, m_messages.cbegin() + first + count, carriesUsage);

    beginRemoveRows(QModelIndex(), first, first + count - 1);
    m_messages.remove(first, count);
    endRemoveRows();

    if (usageChanged)
        emit sessionUsageChanged();
}

QList<MessagePart> ChatModel::processMessageContent(const QString &content) const
{
    QList<MessagePart> parts;
    static const QRegularExpression codeBlockRegex("```(\\w*)\\n?([\\s\\S]*?)```");
    int lastIndex = 0;
    auto blockMatches = codeBlockRegex.globalMatch(content);

    while (blockMatches.hasNext()) {
        auto match = blockMatches.next();
        if (match.capturedStart() > lastIndex) {
            QString textBetween
                = content.mid(lastIndex, match.capturedStart() - lastIndex).trimmed();
            if (!textBetween.isEmpty()) {
                MessagePart part;
                part.type = MessagePartType::Text;
                part.text = textBetween;
                parts.append(part);
            }
        }

        MessagePart codePart;
        codePart.type = MessagePartType::Code;
        codePart.text = match.captured(2).trimmed();
        codePart.language = match.captured(1);
        parts.append(codePart);

        lastIndex = match.capturedEnd();
    }

    if (lastIndex < content.length()) {
        QString remainingText = content.mid(lastIndex).trimmed();

        static const QRegularExpression unclosedBlockRegex("```(\\w*)\\n?([\\s\\S]*)$");
        auto unclosedMatch = unclosedBlockRegex.match(remainingText);

        if (unclosedMatch.hasMatch()) {
            QString beforeCodeBlock = remainingText.left(unclosedMatch.capturedStart()).trimmed();
            if (!beforeCodeBlock.isEmpty()) {
                MessagePart part;
                part.type = MessagePartType::Text;
                part.text = beforeCodeBlock;
                parts.append(part);
            }

            MessagePart codePart;
            codePart.type = MessagePartType::Code;
            codePart.text = unclosedMatch.captured(2).trimmed();
            codePart.language = unclosedMatch.captured(1);
            parts.append(codePart);
        } else if (!remainingText.isEmpty()) {
            MessagePart part;
            part.type = MessagePartType::Text;
            part.text = remainingText;
            parts.append(part);
        }
    }

    return parts;
}

QVariantList ChatModel::userMessagePreviews(int maxLength) const
{
    QVariantList result;
    const int limit = maxLength > 4 ? maxLength : 80;
    for (int i = 0; i < m_messages.size(); ++i) {
        if (m_messages[i].role != ChatRole::User)
            continue;
        QString preview = m_messages[i].content;
        preview.replace(QLatin1Char('\n'), QLatin1Char(' '));
        preview.replace(QLatin1Char('\r'), QLatin1Char(' '));
        preview.replace(QLatin1Char('\t'), QLatin1Char(' '));
        preview = preview.simplified();
        if (preview.size() > limit)
            preview = preview.left(limit - 1).trimmed() + QChar(0x2026);
        QVariantMap entry;
        entry[QStringLiteral("messageIndex")] = i;
        entry[QStringLiteral("preview")] = preview;
        result.append(entry);
    }
    return result;
}

int ChatModel::sessionPromptTokens() const
{
    int total = 0;
    for (const auto &m : m_messages)
        total += m.promptTokens;
    return total;
}

int ChatModel::sessionCompletionTokens() const
{
    int total = 0;
    for (const auto &m : m_messages)
        total += m.completionTokens;
    return total;
}

int ChatModel::sessionCachedPromptTokens() const
{
    int total = 0;
    for (const auto &m : m_messages)
        total += m.cachedPromptTokens;
    return total;
}

int ChatModel::sessionTotalTokens() const
{
    return sessionPromptTokens() + sessionCompletionTokens();
}

void ChatModel::setChatFilePath(const QString &filePath)
{
    m_chatFilePath = filePath;
}

QString ChatModel::chatFilePath() const
{
    return m_chatFilePath;
}

} // namespace QodeAssist::Chat