#include "chatblockmodel.h"

#include <QTextDocument>
#include <QRegularExpressionMatch>


// --- Static helpers ---

static ChatBlock::Type blockTypeFromString(const QString &typeStr)
{
    if (typeStr == "user")           return ChatBlock::User;
    if (typeStr == "assistant")      return ChatBlock::Assistant;
    if (typeStr == "reasoning")      return ChatBlock::Reasoning;
    if (typeStr == "call")           return ChatBlock::Call;
    if (typeStr == "tool")           return ChatBlock::Tool;
    if (typeStr == "error")          return ChatBlock::Error;
    if (typeStr == "info")           return ChatBlock::Info;
    if (typeStr == "retry")          return ChatBlock::Retry;
    return ChatBlock::Default;
}

// --- Block Parsing ---

const QRegularExpression ChatBlockModel::s_headerRe(
    R"(^\[\[\[([^\]]+)\]\]\]$)");

const QRegularExpression ChatBlockModel::s_endRe(
    R"(^\[\[\[end\]\]\]$)");

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

int ChatBlockModel::rowCount(const QModelIndex &parent) const
{
    if (parent.isValid()) return 0;
    return m_blocks.size();
}

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

    const ChatBlock &b = m_blocks[index.row()];
    switch (role) {
	    case BlockTypeRole:
	        switch (b.type) {
	        case ChatBlock::User:      return QStringLiteral("user");
	        case ChatBlock::Assistant: return QStringLiteral("assistant");
	        case ChatBlock::Reasoning: return QStringLiteral("reasoning");
	        case ChatBlock::Call:      return QStringLiteral("call");
	        case ChatBlock::Tool:      return QStringLiteral("tool");
case ChatBlock::Error:     return QStringLiteral("error");
		        case ChatBlock::Info:      return QStringLiteral("info");
		        case ChatBlock::Retry:     return QStringLiteral("retry");
		        case ChatBlock::Interrupt: return QStringLiteral("interrupt");
		        case ChatBlock::Default:   return QStringLiteral("default");
	        }
        return {};
    case BlockNameRole:
        return b.name;
    case ContentRole:
        return b.content;
    case RenderedHtmlRole:
        return b.renderedHtml;
    }
    return {};
}

QHash<int, QByteArray> ChatBlockModel::roleNames() const
{
    return {
        {BlockTypeRole, "blockType"},
        {BlockNameRole, "blockName"},
        {ContentRole,  "blockContent"},
        {RenderedHtmlRole, "blockRenderedHtml"},
    };
}

void ChatBlockModel::appendData(const QString &text)
{
    m_lineBuffer += text;

    int idx;
    while ((idx = m_lineBuffer.indexOf('\n')) >= 0) {
        QString line = m_lineBuffer.left(idx);
        m_lineBuffer = m_lineBuffer.mid(idx + 1);
        processLine(line);
    }
}

void ChatBlockModel::reset()
{
    beginResetModel();
    m_blocks.clear();
    m_lineBuffer.clear();
    m_state = State::BlockStart;
    endResetModel();
    emit countChanged();
}

void ChatBlockModel::trimFront(int keepCount)
{
    if (m_blocks.size() <= keepCount)
        return;
    int removeCount = m_blocks.size() - keepCount;
    beginRemoveRows(QModelIndex(), 0, removeCount - 1);
    m_blocks.remove(0, removeCount);
    endRemoveRows();
    emit countChanged();
}

QString ChatBlockModel::renderMarkdown(const QString &markdown) const
{
    QTextDocument document;
    document.setMarkdown(markdown);
    return document.toHtml();
}

void ChatBlockModel::appendContent(const QString &line)
{
    if (m_blocks.isEmpty())
        return;

    ChatBlock &last = m_blocks.last();
    if (!last.content.isEmpty())
        last.content += '\n';
    last.content += line;
    last.renderedHtml = renderMarkdown(last.content);

    QModelIndex idx = index(m_blocks.size() - 1);
    emit dataChanged(idx, idx, {ContentRole, RenderedHtmlRole});
}

void ChatBlockModel::processLine(const QString &line)
{
    switch (m_state) {
    case State::BlockStart: {
        // [[[blockname]]] is the pre-transition header trigger.
        QRegularExpressionMatch match = s_headerRe.match(line);
        if (match.hasMatch()) {
            const QString blockName = match.captured(1);
            beginInsertRows(QModelIndex(), m_blocks.size(), m_blocks.size());
            m_blocks.append({blockTypeFromString(blockName.section(':', 0, 0)),
                             blockName.section(':', 1), QString(), QString()});
            endInsertRows();
            emit countChanged();
            m_state = State::BlockEnd;
        }
        return;
    }

    case State::BlockEnd:
        // [[[end]]] is the post-transition trigger; all other lines are Markdown.
        if (s_endRe.match(line).hasMatch()) {
            m_state = State::BlockStart;
            return;
        }
        appendContent(line);
        return;
    }
}