#include "chatblockmodel.h"

#include <QRegularExpression>
#include <KSyntaxHighlighting/Repository>
#include <KSyntaxHighlighting/Definition>
#include <KSyntaxHighlighting/Theme>
#include <KSyntaxHighlighting/AbstractHighlighter>
#include <KSyntaxHighlighting/State>
#include <KSyntaxHighlighting/Format>
#include <QGuiApplication>
#include <QClipboard>
#include <QJsonDocument>
#include <QJsonParseError>
#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;
    case FormattedContentRole: {
        // For call blocks: always pretty-print JSON
        // For tool blocks: detect JSON/diff, format accordingly
        // For others: just escape HTML
        QString content = block.content;
        if (content.isEmpty()) return QString();
        if (block.type == ChatBlock::Call) {
            return prettyJson(content);
        }
        QString trimmed = content.trimmed();
        if (trimmed.startsWith('{') || trimmed.startsWith('[')) {
            return prettyJson(content);
        }
        if (trimmed.startsWith(QLatin1String("---")) || trimmed.startsWith(QLatin1String("diff "))) {
            return highlightedHtml(content, QStringLiteral("Diff"));
        }
        // Plain text: escape and add line breaks
        return content.toHtmlEscaped().replace('\n', QLatin1String("<br>"));
    }
    }
    return {};
}

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

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

    // If text is just appended (common during streaming), only reparse the last block
    bool isAppend = text.startsWith(m_rawText) && !m_rawText.isEmpty();
    m_rawText = text;

    if (isAppend && !m_allBlocks.isEmpty()) {
        // Re-parse only: find where the last block started in the text and reparse from there
        // Simpler: just reparse the full text but it's cached via incremental diff
        reparse();
    } else {
        // Full reparse (session switch, initial load)
        m_allBlocks.clear();
        beginResetModel();
        m_blocks.clear();
        m_visibleStart = 0;
        endResetModel();
        reparse();
    }
    emit rawTextChanged();
}

static constexpr int INITIAL_BLOCKS = 20;
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;
}

// Custom highlighter that builds HTML string directly
class HtmlBuilder : 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;
};

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()) {
        def = repo().definitionForFileName("file." + language);
    }
    if (!def.isValid()) {
        QString plain = code.toHtmlEscaped().replace('\n', QStringLiteral("<br>"));
        m_highlightCache.insert(cacheKey, plain);
        return plain;
    }

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

    HtmlBuilder 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 += QStringLiteral("<br>");
    }

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

QString ChatBlockModel::prettyJson(const QString &text) const
{
    // Try to parse and pretty-print JSON, then syntax highlight
    QJsonParseError err;
    QJsonDocument doc = QJsonDocument::fromJson(text.toUtf8(), &err);
    QString formatted;
    if (err.error == QJsonParseError::NoError) {
        formatted = QString::fromUtf8(doc.toJson(QJsonDocument::Indented));
    } else {
        formatted = text;
    }
    return highlightedHtml(formatted, QStringLiteral("JSON"));
}

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