#include "chatblockmodel.h"

#include <QApplication>
#include <QPalette>
#include <QRegularExpressionMatch>
#include <KSyntaxHighlighting/Repository>
#include <KSyntaxHighlighting/Definition>
#include <KSyntaxHighlighting/Theme>
#include <KSyntaxHighlighting/AbstractHighlighter>
#include <KSyntaxHighlighting/Format>
#include <KSyntaxHighlighting/State>

// --- Syntax Highlighting (copied from Kate plugin ChatRenderer) ---

class HtmlHighlighter : public KSyntaxHighlighting::AbstractHighlighter
{
public:
    QString html;

    using AbstractHighlighter::AbstractHighlighter;

    KSyntaxHighlighting::State processLine(const QString &line, const KSyntaxHighlighting::State &state)
    {
        m_currentLine = line;
        return highlightLine(line, state);
    }

protected:
    void applyFormat(int offset, int length, const KSyntaxHighlighting::Format &format) override
    {
        QString text = m_currentLine.mid(offset, length).toHtmlEscaped();
        if (format.hasTextColor(theme())) {
            QColor fg = format.textColor(theme());
            html += QStringLiteral("<span style=\"color:") + fg.name() + QStringLiteral("\">") + text + QStringLiteral("</span>");
        } else {
            html += text;
        }
    }

private:
    QString m_currentLine;
};

// --- Block Parsing ---

const QRegularExpression ChatBlockModel::s_headerRe(
    R"(^\[(user|assistant|reasoning|call|tool|error|info)(?::([^\]]*))?\]$)");

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::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();
    // Flush any unclosed fence buffer as plain prose
    m_fenceState = NoFence;
    m_fenceLang.clear();
    m_fenceBuffer.clear();
    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::escapeHtml(const QString &text) const
{
    return text.toHtmlEscaped();
}

QString ChatBlockModel::highlightCode(const QString &code, const QString &language) const
{
    KSyntaxHighlighting::Repository *repo = new KSyntaxHighlighting::Repository();
    KSyntaxHighlighting::Definition def = repo->definitionForName(language);
    if (!def.isValid())
        def = repo->definitionForFileName("file." + language);

    bool dark = QApplication::palette().color(QPalette::Window).lightnessF() < 0.5;
    KSyntaxHighlighting::Theme theme = repo->defaultTheme(dark
        ? KSyntaxHighlighting::Repository::DarkTheme
        : KSyntaxHighlighting::Repository::LightTheme);

    if (!def.isValid()) {
        delete repo;
        return QStringLiteral("<pre style=\"white-space:pre-wrap;\">") + escapeHtml(code) + QStringLiteral("</pre>");
    }

    HtmlHighlighter builder;
    builder.setDefinition(def);
    builder.setTheme(theme);

    KSyntaxHighlighting::State state;
    const QStringList lines = code.split('\n');
    for (int i = 0; i < lines.size(); ++i) {
        state = builder.processLine(lines[i], state);
        if (i < lines.size() - 1)
            builder.html += '\n';
    }

    delete repo;
    return QStringLiteral("<pre style=\"white-space:pre-wrap;\">") + builder.html + QStringLiteral("</pre>");
}

QString ChatBlockModel::renderDiff(const QString &diff) const
{
    QString html = QStringLiteral("<pre style=\"white-space:pre-wrap;\">");
    const QStringList lines = diff.split('\n');
    for (const QString &line : lines) {
        if (line.startsWith('+')) {
            html += QStringLiteral("<span style=\"color:#4caf50;\">") + escapeHtml(line) + QStringLiteral("</span>\n");
        } else if (line.startsWith('-')) {
            html += QStringLiteral("<span style=\"color:#f44336;\">") + escapeHtml(line) + QStringLiteral("</span>\n");
        } else if (line.startsWith("@@")) {
            html += QStringLiteral("<span style=\"color:#2196f3;\">") + escapeHtml(line) + QStringLiteral("</span>\n");
        } else {
            html += escapeHtml(line) + '\n';
        }
    }
    html += QStringLiteral("</pre>");
    return html;
}

void ChatBlockModel::processLine(const QString &line)
{
    // --- Fence state machine ---
    if (m_fenceState == InFence) {
        // Check for closing fence: ^```\s*$
        if (line.trimmed() == "```") {
            // Close fence — render and finalize block
            QString rendered;
            if (m_fenceLang == "diff" || (m_fenceLang.isEmpty() &&
                (m_fenceBuffer.startsWith("--- ") || m_fenceBuffer.startsWith("diff --git")))) {
                rendered = renderDiff(m_fenceBuffer);
            } else {
                rendered = highlightCode(m_fenceBuffer, m_fenceLang);
            }

            if (!m_blocks.isEmpty()) {
                ChatBlock &last = m_blocks.last();
                last.renderedHtml = rendered;
                QModelIndex idx = index(m_blocks.size() - 1);
                emit dataChanged(idx, idx, {RenderedHtmlRole});
            }

            m_fenceState = NoFence;
            m_fenceLang.clear();
            m_fenceBuffer.clear();
            return;
        }

        // Accumulate fenced content
        if (!m_fenceBuffer.isEmpty())
            m_fenceBuffer += '\n';
        m_fenceBuffer += line;
        return;
    }

    // --- NoFence state: look for opening fence or header ---
    static const QRegularExpression fenceOpenRe(R"(^```(\w*)[ \t]*$)");
    QRegularExpressionMatch fenceMatch = fenceOpenRe.match(line);
    if (fenceMatch.hasMatch()) {
        // Open fence — start buffering. Create a placeholder block if needed.
        m_fenceState = InFence;
        m_fenceLang = fenceMatch.captured(1);
        m_fenceBuffer.clear();

        // Ensure there's a block to attach the rendered result to
        if (m_blocks.isEmpty()) {
            beginInsertRows(QModelIndex(), 0, 0);
            m_blocks.append({ChatBlock::Default, QString(), QString(), QString()});
            endInsertRows();
            emit countChanged();
        }
        return;
    }

    // Header line
    QRegularExpressionMatch m = s_headerRe.match(line);
    if (m.hasMatch()) {
        QString typeStr = m.captured(1);
        QString qualifier = m.captured(2);

        ChatBlock::Type type;
        if (typeStr == "user")           type = ChatBlock::User;
        else if (typeStr == "assistant") type = ChatBlock::Assistant;
        else if (typeStr == "reasoning") type = ChatBlock::Reasoning;
        else if (typeStr == "call")      type = ChatBlock::Call;
        else if (typeStr == "tool")      type = ChatBlock::Tool;
        else if (typeStr == "error")     type = ChatBlock::Error;
        else if (typeStr == "info")      type = ChatBlock::Info;
        else                             type = ChatBlock::Default;

        beginInsertRows(QModelIndex(), m_blocks.size(), m_blocks.size());
        m_blocks.append({type, qualifier, QString(), QString()});
        endInsertRows();
        emit countChanged();
        return;
    }

    // Content line — append to the last block if any
    if (m_blocks.isEmpty()) {
        // Mid-stream connect: create a Default block
        beginInsertRows(QModelIndex(), 0, 0);
        m_blocks.append({ChatBlock::Default, QString(), QString(), QString()});
        endInsertRows();
        emit countChanged();
    }

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

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