#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;
    QString defaultColor;  // Fallback color for unhighlighted text
    using AbstractHighlighter::AbstractHighlighter;

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

    void endLine() {
        // Emit any remaining unhighlighted text at end of line
        if (m_offset < m_line.length()) {
            html += escapeHtml(m_line.mid(m_offset));
        }
    }

protected:
    void applyFormat(int offset, int length,
                     const KSyntaxHighlighting::Format &format) override
    {
        // Emit any unhighlighted text before this span
        if (offset > m_offset) {
            const QString gap = m_line.mid(m_offset, offset - m_offset);
            html += escapeHtml(gap);
        }
        m_offset = offset + length;

        const QString text = escapeHtml(m_line.mid(offset, length));
        if (format.hasTextColor(theme())) {
            html += QStringLiteral("<span style=\"color:")
                + format.textColor(theme()).name()
                + QStringLiteral("\">") + text + QStringLiteral("</span>");
        } else if (!defaultColor.isEmpty()) {
            html += QStringLiteral("<span style=\"color:")
                + defaultColor
                + QStringLiteral("\">") + text + QStringLiteral("</span>");
        } else {
            html += text;
        }
    }

private:
    QString escapeHtml(const QString &text) const {
        return text.toHtmlEscaped()
            .replace(QStringLiteral(" "), QStringLiteral("&nbsp;"))
            .replace(QStringLiteral("\t"), QStringLiteral("&nbsp;&nbsp;&nbsp;&nbsp;"));
    }

    QString m_line;
    int m_offset = 0;
};

// --- 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\]\]\]$)");

// Matches opening fence: ```lang or just ```
const QRegularExpression ChatBlockModel::s_fenceOpenRe(
    R"(^```([A-Za-z0-9_+.-]*).*$)");

// Matches closing fence: ``` on its own line
const QRegularExpression ChatBlockModel::s_fenceCloseRe(
    R"(^```$)");

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

    // For streaming: update current block with partial line buffer content
    // so the UI shows text as it arrives, not just after newlines.
    if (!m_lineBuffer.isEmpty() && !m_blocks.isEmpty() && m_state == State::BlockEnd) {
        ChatBlock &last = m_blocks.last();
        // Store the committed content length so we can update partial content
        QString fullContent = last.content;
        if (!fullContent.isEmpty() && !m_lineBuffer.isEmpty())
            fullContent += '\n';
        fullContent += m_lineBuffer;
        
        // Temporarily set content with partial line for rendering
        QString savedContent = last.content;
        last.content = fullContent;
        rebuildBlockHtml(m_blocks.size() - 1);
        last.content = savedContent;  // Restore — partial line not committed yet
        
        // But emit with the rendered HTML that includes partial content
        QModelIndex idx = index(m_blocks.size() - 1);
        emit dataChanged(idx, idx, {RenderedHtmlRole});
    }
}

void ChatBlockModel::reset()
{
    beginResetModel();
    m_blocks.clear();
    m_lineBuffer.clear();
    m_state = State::BlockStart;
    m_inCodeFence = false;
    m_fenceLanguage.clear();
    m_fenceContent.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);
#if QT_VERSION >= QT_VERSION_CHECK(6, 0, 0)
    m_blocks.remove(0, removeCount);
#else
    m_blocks.erase(m_blocks.begin(), m_blocks.begin() + removeCount);
#endif
    endRemoveRows();
    emit countChanged();
}

QString ChatBlockModel::getDefaultTextColor() const
{
    const auto textColor = QApplication::palette().color(QPalette::Text);
    return textColor.name();
}

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

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

    const QString defaultColor = getDefaultTextColor();

    // No valid definition — just escape and return
    if (!definition.isValid()) {
        QString result = QStringLiteral("<pre style=\"white-space:pre-wrap;color:") + defaultColor + QStringLiteral("\">");
        result += code.toHtmlEscaped();
        result += QStringLiteral("</pre>");
        return result;
    }

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

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

    QString result = QStringLiteral("<pre style=\"white-space:pre-wrap;\">");
    result += highlighter.html;
    result += QStringLiteral("</pre>");
    return result;
}

QString ChatBlockModel::escapeForHtml(const QString &text) const
{
    const QString defaultColor = getDefaultTextColor();
    QString result = QStringLiteral("<div style=\"white-space:pre-wrap;color:") + defaultColor + QStringLiteral("\">");
    result += text.toHtmlEscaped();
    result += QStringLiteral("</div>");
    return result;
}

void ChatBlockModel::rebuildBlockHtml(int blockIndex)
{
    if (blockIndex < 0 || blockIndex >= m_blocks.size())
        return;

    ChatBlock &block = m_blocks[blockIndex];
    QString html;

    // Parse the block content for code fences
    const QStringList lines = block.content.split('\n');
    bool inFence = false;
    QString fenceLang;
    QStringList fenceLines;
    QStringList proseLines;

    auto flushProse = [&]() {
        if (!proseLines.isEmpty()) {
            html += escapeForHtml(proseLines.join('\n'));
            proseLines.clear();
        }
    };

    auto flushFence = [&]() {
        if (!fenceLines.isEmpty()) {
            html += highlightCode(fenceLines.join('\n'), fenceLang);
            fenceLines.clear();
            fenceLang.clear();
        }
    };

    for (const QString &line : lines) {
        if (!inFence) {
            // Check for fence open
            auto match = s_fenceOpenRe.match(line);
            if (match.hasMatch()) {
                flushProse();
                inFence = true;
                fenceLang = match.captured(1);
                continue;
            }
            proseLines.append(line);
        } else {
            // Inside fence — check for close
            if (s_fenceCloseRe.match(line).hasMatch()) {
                flushFence();
                inFence = false;
                continue;
            }
            fenceLines.append(line);
        }
    }

    // Handle unclosed fence (streaming) or remaining prose
    if (inFence) {
        // Still inside a fence — render what we have
        html += highlightCode(fenceLines.join('\n'), fenceLang);
    } else {
        flushProse();
    }

    block.renderedHtml = html;
}

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;

    rebuildBlockHtml(m_blocks.size() - 1);

    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]]] starts a new block
        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;
        }
        // Lines outside blocks are ignored
        return;
    }

    case State::BlockEnd:
        // [[[end]]] closes the current block
        if (s_endRe.match(line).hasMatch()) {
            m_state = State::BlockStart;
            return;
        }
        appendContent(line);
        return;
    }
}