#include "chatblockmodel.h"

#include <QTextDocument>
#include <QRegularExpressionMatch>
#include <QApplication>
#include <QPalette>

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


// --- Syntax highlighting ---

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

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

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

private:
    QString m_line;
};

// --- 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)
    , m_repo(new KSyntaxHighlighting::Repository())
{
}

ChatBlockModel::~ChatBlockModel()
{
    delete m_repo;
}

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::highlightCode(const QString &code, const QString &language) const
{
    auto definition = m_repo->definitionForName(language);
    if (!definition.isValid())
        definition = m_repo->definitionForFileName(QStringLiteral("file.") + language);
    if (!definition.isValid())
        return QStringLiteral("<pre style=\"white-space:pre-wrap;\">")
            + code.toHtmlEscaped() + QStringLiteral("</pre>");

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

    HtmlHighlighter highlighter;
    highlighter.setDefinition(definition);
    highlighter.setTheme(theme);

    KSyntaxHighlighting::State state;
    const QStringList lines = code.split('\n');
    for (int i = 0; i < lines.size(); ++i) {
        state = highlighter.processLine(lines[i], state);
        if (i + 1 < lines.size())
            highlighter.html += QStringLiteral("<br/>");
    }
    return QStringLiteral("<pre style=\"white-space:pre-wrap;\">")
        + highlighter.html + QStringLiteral("</pre>");
}

QString ChatBlockModel::renderMarkdown(const QString &markdown) const
{
    static const QRegularExpression fenceRe(
        R"(```([A-Za-z0-9_+.-]*)[^\n]*\n([\s\S]*?)```)");

    QString result;
    int cursor = 0;
    auto it = fenceRe.globalMatch(markdown);
    while (it.hasNext()) {
        const auto match = it.next();
        // Prose before this fence: HTML-escape and wrap in pre-wrap div to
        // preserve newlines and whitespace. QTextDocument::setMarkdown would
        // collapse single newlines (standard Markdown behavior).
        const QString prose = markdown.mid(cursor, match.capturedStart() - cursor);
        if (!prose.isEmpty())
            result += QStringLiteral("<div style=\"white-space:pre-wrap\">")
                + prose.toHtmlEscaped() + QStringLiteral("</div>");
        result += highlightCode(match.captured(2), match.captured(1));
        cursor = match.capturedEnd();
    }

    // During streaming, the closing fence may not have arrived yet. Do not
    // send the incomplete code block through Markdown, which collapses its
    // newlines into prose spacing.
    static const QRegularExpression openFenceRe(
        R"(^```([A-Za-z0-9_+.-]*)[^\n]*\n([\s\S]*)$)",
        QRegularExpression::MultilineOption);
    const auto tail = markdown.mid(cursor);
    if (!tail.isEmpty()) {
        const auto open = openFenceRe.match(tail);
        if (open.hasMatch()) {
            const QString prose = tail.left(open.capturedStart());
            if (!prose.isEmpty())
                result += QStringLiteral("<div style=\"white-space:pre-wrap\">")
                    + prose.toHtmlEscaped() + QStringLiteral("</div>");
            result += highlightCode(open.captured(2), open.captured(1));
        } else {
            result += QStringLiteral("<div style=\"white-space:pre-wrap\">")
                + tail.toHtmlEscaped() + QStringLiteral("</div>");
        }
    }
    return result;
}

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;
    }
}