#include "chatblockmodel.h"

#include <QRegularExpression>
#include <KSyntaxHighlighting/Repository>
#include <KSyntaxHighlighting/Definition>
#include <KSyntaxHighlighting/Theme>
#include <KSyntaxHighlighting/SyntaxHighlighter>
#include <QTextDocument>
#include <QGuiApplication>
#include <QClipboard>
#include <QPalette>

static KSyntaxHighlighting::Repository *s_repo = nullptr;

static KSyntaxHighlighting::Repository &repo()
{
    if (!s_repo)
        s_repo = new KSyntaxHighlighting::Repository();
    return *s_repo;
}

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 &block = m_blocks[index.row()];
    switch (role) {
    case BlockTypeRole: {
        switch (block.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");
        }
        return {};
    }
    case NameRole:
        return block.name;
    case ContentRole:
        return block.content;
    case SubBlocksRole:
        if (block.cachedContent != block.content) {
            block.cachedSubBlocks = parseSubBlocks(block.content);
            block.cachedContent = block.content;
        }
        return block.cachedSubBlocks;
    }
    return {};
}

QHash<int, QByteArray> ChatBlockModel::roleNames() const
{
    return {
        {BlockTypeRole, "blockType"},
        {NameRole, "blockName"},
        {ContentRole, "blockContent"},
        {SubBlocksRole, "subBlocks"},
    };
}

void ChatBlockModel::setRawText(const QString &text)
{
    if (text == m_rawText) return;
    m_rawText = text;
    reparse();
    emit rawTextChanged();
}

static constexpr int INITIAL_BLOCKS = 40;
static constexpr int LOAD_MORE_COUNT = 30;

void ChatBlockModel::reparse()
{
    m_allBlocks = parseBlocks(m_rawText);

    // Determine visible window: last INITIAL_BLOCKS (or expand if already showing more)
    int desiredStart;
    if (m_blocks.isEmpty()) {
        // Initial load: show last N blocks
        desiredStart = qMax(0, m_allBlocks.size() - INITIAL_BLOCKS);
    } else {
        // Streaming update: keep the same start offset, just extend the tail
        desiredStart = m_visibleStart;
        // But if allBlocks shrank (session switch), reset
        if (desiredStart >= m_allBlocks.size())
            desiredStart = qMax(0, m_allBlocks.size() - INITIAL_BLOCKS);
    }
    m_visibleStart = desiredStart;

    QList<ChatBlock> newBlocks = m_allBlocks.mid(m_visibleStart);

    // Incremental update: only change what differs
    int commonCount = qMin(m_blocks.size(), newBlocks.size());
    int firstDiff = commonCount;

    for (int i = 0; i < commonCount; ++i) {
        if (m_blocks[i].content != newBlocks[i].content ||
            m_blocks[i].type != newBlocks[i].type ||
            m_blocks[i].name != newBlocks[i].name) {
            firstDiff = i;
            break;
        }
    }

    // Remove extra old blocks
    if (m_blocks.size() > newBlocks.size()) {
        beginRemoveRows(QModelIndex(), newBlocks.size(), m_blocks.size() - 1);
        m_blocks = m_blocks.mid(0, newBlocks.size());
        endRemoveRows();
    }

    // Update changed existing blocks
    for (int i = firstDiff; i < qMin(m_blocks.size(), newBlocks.size()); ++i) {
        m_blocks[i] = newBlocks[i];
        emit dataChanged(index(i), index(i));
    }

    // Append new blocks at end
    if (newBlocks.size() > m_blocks.size()) {
        beginInsertRows(QModelIndex(), m_blocks.size(), newBlocks.size() - 1);
        for (int i = m_blocks.size(); i < newBlocks.size(); ++i) {
            m_blocks.append(newBlocks[i]);
        }
        endInsertRows();
    }
}

bool ChatBlockModel::loadEarlier()
{
    if (m_visibleStart <= 0)
        return false; // already showing everything

    int loadCount = qMin(LOAD_MORE_COUNT, m_visibleStart);
    int newStart = m_visibleStart - loadCount;

    QList<ChatBlock> prepend = m_allBlocks.mid(newStart, loadCount);

    beginInsertRows(QModelIndex(), 0, loadCount - 1);
    for (int i = loadCount - 1; i >= 0; --i) {
        m_blocks.prepend(prepend[i]);
    }
    m_visibleStart = newStart;
    endInsertRows();

    return true;
}

QList<ChatBlock> ChatBlockModel::parseBlocks(const QString &text) const
{
    QList<ChatBlock> blocks;
    if (text.isEmpty()) return blocks;

    // Block headers: [user], [assistant], [reasoning], [call:name], [tool:name], [error], [info] msg
    static const QRegularExpression headerRe(
        R"(^\[(user|assistant|reasoning|call|tool|error|info)(?::([^\]]*))?\](.*)$)",
        QRegularExpression::MultilineOption);

    auto it = headerRe.globalMatch(text);
    QList<QPair<int, QRegularExpressionMatch>> headers;
    while (it.hasNext()) {
        auto match = it.next();
        headers.append({match.capturedStart(), match});
    }

    if (headers.isEmpty()) {
        // No headers found — treat entire text as assistant block
        ChatBlock b;
        b.type = ChatBlock::Assistant;
        b.content = text;
        blocks.append(b);
        return blocks;
    }

    for (int i = 0; i < headers.size(); ++i) {
        const auto &[pos, match] = headers[i];
        QString typeStr = match.captured(1);
        QString name = match.captured(2);
        QString inlineContent = match.captured(3); // for [info] msg style

        int contentStart = match.capturedEnd();
        int contentEnd = (i + 1 < headers.size()) ? headers[i + 1].first : text.size();

        QString content;
        if (typeStr == "info" && !inlineContent.trimmed().isEmpty()) {
            // [info] has inline content on same line
            content = inlineContent.trimmed();
        } else {
            content = text.mid(contentStart, contentEnd - contentStart);
            // Strip leading newline and trailing double-newline
            if (content.startsWith('\n'))
                content = content.mid(1);
            while (content.endsWith("\n\n"))
                content.chop(1);
        }

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

        b.name = name;
        b.content = content;
        blocks.append(b);
    }

    return blocks;
}

QVariantList ChatBlockModel::parseSubBlocks(const QString &content) const
{
    QVariantList result;
    if (content.isEmpty()) return result;

    // Split on code fences: ```language\n...\n```
    static const QRegularExpression fenceRe(
        R"(^```(\w*)\s*\n((?:(?!^```).)*?)\n^```\s*$)",
        QRegularExpression::MultilineOption | QRegularExpression::DotMatchesEverythingOption);

    int lastEnd = 0;
    auto it = fenceRe.globalMatch(content);
    while (it.hasNext()) {
        auto match = it.next();
        // Prose before this fence
        if (match.capturedStart() > lastEnd) {
            QString prose = content.mid(lastEnd, match.capturedStart() - lastEnd).trimmed();
            if (!prose.isEmpty()) {
                QVariantMap m;
                m["type"] = "prose";
                m["content"] = prose;
                result.append(m);
            }
        }

        QString lang = match.captured(1);
        QString code = match.captured(2);

        // Detect diffs
        QString type = "code";
        if (lang == "diff" || (lang.isEmpty() &&
            (code.startsWith("--- ") || code.startsWith("diff --git")))) {
            type = "diff";
        }

        QVariantMap m;
        m["type"] = type;
        m["language"] = lang;
        m["content"] = code;
        result.append(m);

        lastEnd = match.capturedEnd();
    }

    // Trailing prose
    if (lastEnd < content.size()) {
        QString prose = content.mid(lastEnd).trimmed();
        if (!prose.isEmpty()) {
            QVariantMap m;
            m["type"] = "prose";
            m["content"] = prose;
            result.append(m);
        }
    }

    // If no fences found at all, the entire content is prose
    if (result.isEmpty()) {
        QVariantMap m;
        m["type"] = "prose";
        m["content"] = content;
        result.append(m);
    }

    return result;
}

QString ChatBlockModel::highlightedHtml(const QString &code, const QString &language) const
{
    // Check cache first
    QString cacheKey = language + "\x1f" + code;
    auto cached = m_highlightCache.find(cacheKey);
    if (cached != m_highlightCache.end())
        return *cached;

    auto def = repo().definitionForName(language);
    if (!def.isValid()) {
        // Try by file extension
        def = repo().definitionForFileName("file." + language);
    }
    if (!def.isValid()) {
        // Return escaped plain text
        QString plain = code.toHtmlEscaped().replace('\n', "<br>");
        m_highlightCache.insert(cacheKey, plain);
        return plain;
    }

    // Determine theme based on current palette
    bool dark = QGuiApplication::palette().color(QPalette::Window).lightnessF() < 0.5;
    auto theme = dark
        ? repo().defaultTheme(KSyntaxHighlighting::Repository::DarkTheme)
        : repo().defaultTheme(KSyntaxHighlighting::Repository::LightTheme);

    QTextDocument doc;
    doc.setPlainText(code);

    auto *highlighter = new KSyntaxHighlighting::SyntaxHighlighter(&doc);
    highlighter->setDefinition(def);
    highlighter->setTheme(theme);
    highlighter->rehighlight();

    // Extract HTML from the highlighted document
    QString html;
    QTextBlock block = doc.begin();
    while (block.isValid()) {
        for (auto it = block.begin(); !it.atEnd(); ++it) {
            QTextFragment fragment = it.fragment();
            if (!fragment.isValid()) continue;

            QTextCharFormat fmt = fragment.charFormat();
            QString text = fragment.text().toHtmlEscaped();

            if (fmt.foreground().style() != Qt::NoBrush) {
                QColor fg = fmt.foreground().color();
                html += QString("<span style=\"color:%1;\">%2</span>").arg(fg.name(), text);
            } else {
                html += text;
            }
        }
        block = block.next();
        if (block.isValid())
            html += "<br>";
    }

    m_highlightCache.insert(cacheKey, html);
    return html;
}

void ChatBlockModel::copyToClipboard(const QString &text) const
{
    QGuiApplication::clipboard()->setText(text);
}