#include "chat_renderer.h"

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

#include <QApplication>
#include <QPalette>
#include <QTextCursor>
#include <QScrollBar>
#include <QRegularExpression>

// --- HTML Highlighter ---

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

// --- ChatRenderer ---

ChatRenderer::ChatRenderer()
    : m_repo(new KSyntaxHighlighting::Repository())
{
}

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

void ChatRenderer::detectTheme() const
{
    m_dark = QApplication::palette().color(QPalette::Window).lightnessF() < 0.5;
}

QString ChatRenderer::escapeHtml(const QString &text) const
{
    return text.toHtmlEscaped();
}

// --- Block Parsing ---

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

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

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

        QString content;
        if (typeStr == "info" && !inlineContent.trimmed().isEmpty()) {
            content = inlineContent.trimmed();
        } else {
            content = text.mid(contentStart, contentEnd - contentStart);
            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;
}

QList<SubBlock> ChatRenderer::parseSubBlocks(const QString &content) const
{
    QList<SubBlock> result;
    if (content.isEmpty()) return result;

    static const QRegularExpression interruptRe(
        R"(<system-user-interruption>\n?(.*?)\n?</system-user-interruption>)",
        QRegularExpression::DotMatchesEverythingOption);

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

    // Find all interrupt tags
    int lastEnd = 0;
    auto intIt = interruptRe.globalMatch(content);
    QList<QPair<int, QRegularExpressionMatch>> interrupts;
    while (intIt.hasNext()) {
        auto match = intIt.next();
        interrupts.append({match.capturedStart(), match});
    }

    // Helper: parse a segment for code fences and prose
    auto parseSegment = [&](const QString &segment) {
        if (segment.trimmed().isEmpty()) return;

        int segLastEnd = 0;
        auto it = fenceRe.globalMatch(segment);
        while (it.hasNext()) {
            auto match = it.next();
            if (match.capturedStart() > segLastEnd) {
                QString prose = segment.mid(segLastEnd, match.capturedStart() - segLastEnd).trimmed();
                if (!prose.isEmpty()) {
                    SubBlock sb;
                    sb.type = SubBlock::Prose;
                    sb.content = prose;
                    result.append(sb);
                }
            }

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

            SubBlock sb;
            if (lang == "diff" || (lang.isEmpty() &&
                (code.startsWith("--- ") || code.startsWith("diff --git")))) {
                sb.type = SubBlock::Diff;
            } else {
                sb.type = SubBlock::Code;
            }
            sb.language = lang;
            sb.content = code;
            result.append(sb);

            segLastEnd = match.capturedEnd();
        }

        if (segLastEnd < segment.size()) {
            QString prose = segment.mid(segLastEnd).trimmed();
            if (!prose.isEmpty()) {
                SubBlock sb;
                sb.type = SubBlock::Prose;
                sb.content = prose;
                result.append(sb);
            }
        }
    };

    if (interrupts.isEmpty()) {
        parseSegment(content);
    } else {
        for (int i = 0; i < interrupts.size(); ++i) {
            const auto &[pos, match] = interrupts[i];

            if (pos > lastEnd) {
                parseSegment(content.mid(lastEnd, pos - lastEnd));
            }

            SubBlock sb;
            sb.type = SubBlock::Interrupt;
            sb.content = match.captured(1).trimmed();
            result.append(sb);

            lastEnd = match.capturedEnd();
        }

        if (lastEnd < content.size()) {
            parseSegment(content.mid(lastEnd));
        }
    }

    if (result.isEmpty()) {
        SubBlock sb;
        sb.type = SubBlock::Prose;
        sb.content = content;
        result.append(sb);
    }

    return result;
}

// --- HTML Rendering ---

QString ChatRenderer::highlightCode(const QString &code, const QString &language) const
{
    auto def = m_repo->definitionForName(language);
    if (!def.isValid())
        def = m_repo->definitionForFileName("file." + language);
    if (!def.isValid())
        return QStringLiteral("<pre style=\"margin:4px 0;white-space:pre-wrap;\">") + escapeHtml(code) + QStringLiteral("</pre>");

    auto theme = m_dark
        ? m_repo->defaultTheme(KSyntaxHighlighting::Repository::DarkTheme)
        : m_repo->defaultTheme(KSyntaxHighlighting::Repository::LightTheme);

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

    return QStringLiteral("<pre style=\"margin:4px 0;white-space:pre-wrap;\">") + builder.html + QStringLiteral("</pre>");
}

QString ChatRenderer::renderDiff(const QString &diff) const
{
    QString html = QStringLiteral("<pre style=\"margin:4px 0;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;
}

QString ChatRenderer::renderProse(const QString &prose) const
{
    // Simple markdown-ish rendering: bold, italic, inline code, links
    QString html = escapeHtml(prose);

    // Inline code: `code`
    static const QRegularExpression inlineCode("`([^`]+)`");
    html.replace(inlineCode, QStringLiteral("<code style=\"background:") +
        (m_dark ? QStringLiteral("#333") : QStringLiteral("#f0f0f0")) +
        QStringLiteral(";padding:1px 4px;border-radius:3px;\">\\1</code>"));

    // Bold: **text** or __text__
    static const QRegularExpression bold(R"(\*\*(.+?)\*\*)");
    html.replace(bold, QStringLiteral("<b>\\1</b>"));
    static const QRegularExpression bold2(R"(__(.+?)__)");
    html.replace(bold2, QStringLiteral("<b>\\1</b>"));

    // Italic: *text* or _text_
    static const QRegularExpression italic(R"((?<!\*)\*(?!\*)(.+?)(?<!\*)\*(?!\*))");
    html.replace(italic, QStringLiteral("<i>\\1</i>"));

    // Bullet points: lines starting with - or *
    static const QRegularExpression bullet(R"(^(\s*)[\-\*]\s)", QRegularExpression::MultilineOption);
    html.replace(bullet, QStringLiteral("\\1• "));

    // Newlines to <br>
    html.replace('\n', QStringLiteral("<br>"));

    return html;
}

QString ChatRenderer::subBlockToHtml(const SubBlock &sub, bool preserveWhitespace) const
{
    QString codeBg = m_dark ? QStringLiteral("#1e1e2e") : QStringLiteral("#f5f5f5");
    QString codeBorder = m_dark ? QStringLiteral("#444") : QStringLiteral("#ddd");

    switch (sub.type) {
    case SubBlock::Code: {
        QString langLabel;
        if (!sub.language.isEmpty()) {
            langLabel = QStringLiteral(
                "<div style=\"font-size:small;color:gray;padding:2px 8px;\">") +
                escapeHtml(sub.language) + QStringLiteral("</div>");
        }
        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:6px 0;\">"
            "<tr><td style=\"background:") + codeBg +
            QStringLiteral(";border:1px solid ") + codeBorder +
            QStringLiteral(";padding:8px;\">") +
            langLabel + highlightCode(sub.content, sub.language) +
            QStringLiteral("</td></tr></table>");
    }
    case SubBlock::Diff: {
        // Store diff and render a clickable link to open in Kate's diff viewer
        if (m_diffStorage) {
            int idx = m_diffStorage->size();
            m_diffStorage->append(sub.content);

            // Count lines for display
            int adds = 0, dels = 0;
            const QStringList lines = sub.content.split('\n');
            for (const QString &line : lines) {
                if (line.startsWith('+') && !line.startsWith("+++")) adds++;
                else if (line.startsWith('-') && !line.startsWith("---")) dels++;
            }
            QString stats;
            if (adds > 0) stats += QStringLiteral("<span style=\"color:#4caf50;\">+%1</span> ").arg(adds);
            if (dels > 0) stats += QStringLiteral("<span style=\"color:#f44336;\">-%1</span>").arg(dels);

            return QStringLiteral(
                "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:6px 0;\">"
                "<tr><td style=\"background:") + codeBg +
                QStringLiteral(";border:1px solid ") + codeBorder +
                QStringLiteral(";padding:8px 12px;\">"
                "<span style=\"font-size:small;color:gray;\">▶ diff</span> ") +
                stats +
                QStringLiteral("<br><a href=\"ollie-diff:///") + QString::number(idx) +
                QStringLiteral("\" style=\"color:#2196f3;text-decoration:underline;\">"
                "↗ Open in Diff Viewer</a>"
                "</td></tr></table>");
        }

        // Fallback: render inline
        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:6px 0;\">"
            "<tr><td style=\"background:") + codeBg +
            QStringLiteral(";border:1px solid ") + codeBorder +
            QStringLiteral(";padding:8px;\">"
            "<div style=\"font-size:small;color:gray;padding:2px 8px;\">diff</div>") +
            renderDiff(sub.content) +
            QStringLiteral("</td></tr></table>");
    }
    case SubBlock::Prose:
        if (preserveWhitespace) {
            return QStringLiteral("<pre style=\"margin:4px 0;white-space:pre-wrap;font-family:monospace;font-size:small;\">") +
                escapeHtml(sub.content) + QStringLiteral("</pre>");
        }
        return QStringLiteral("<div style=\"margin:4px 0;line-height:1.4;\">") +
            renderProse(sub.content) + QStringLiteral("</div>");
    case SubBlock::Interrupt: {
        QString interruptBg = m_dark ? QStringLiteral("#3d2e00") : QStringLiteral("#fff8e1");
        QString interruptBorder = QStringLiteral("#ffc107");
        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:8px 0;\">"
            "<tr><td style=\"background:") + interruptBg +
            QStringLiteral(";padding:8px 12px;border-left:3px solid ") + interruptBorder +
            QStringLiteral(";\">"
            "<div style=\"font-size:small;color:") + interruptBorder +
            QStringLiteral(";font-weight:bold;\">⚠ user interrupt</div>"
            "<div style=\"margin:4px 0;line-height:1.4;\">") +
            renderProse(sub.content) +
            QStringLiteral("</div></td></tr></table>");
    }
    }
    return {};
}

QString ChatRenderer::blockToHtml(const ChatBlock &block) const
{
    detectTheme();

    QString userBg = m_dark ? QStringLiteral("#1a3a5c") : QStringLiteral("#e3f2fd");
    QString assistantBg = m_dark ? QStringLiteral("#2a2a2a") : QStringLiteral("#fafafa");
    QString errorBg = m_dark ? QStringLiteral("#3d1515") : QStringLiteral("#ffebee");
    QString metaBg = m_dark ? QStringLiteral("#1a1a2e") : QStringLiteral("#f5f5f5");
    QString metaColor = m_dark ? QStringLiteral("#888") : QStringLiteral("#666");
    QString borderColor = m_dark ? QStringLiteral("#444") : QStringLiteral("#e0e0e0");

    switch (block.type) {
    case ChatBlock::User: {
        QList<SubBlock> subs = parseSubBlocks(block.content);
        QString inner;
        for (const auto &sub : subs)
            inner += subBlockToHtml(sub);

        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:8px 0;\">"
            "<tr><td width=\"15%\"></td>"
            "<td style=\"background:") + userBg +
            QStringLiteral(";padding:10px 14px;border:1px solid ") + borderColor +
            QStringLiteral(";\">") + inner +
            QStringLiteral("</td></tr>"
            "<tr><td></td><td align=\"right\"><span style=\"font-size:small;color:gray;\">You</span></td></tr>"
            "</table>");
    }
    case ChatBlock::Assistant: {
        QList<SubBlock> subs = parseSubBlocks(block.content);
        QString inner;
        QString lastProse; // track preceding prose for file path extraction
        for (const auto &sub : subs) {
            if (sub.type == SubBlock::Prose) {
                lastProse = sub.content;
            }
            inner += subBlockToHtml(sub);

            // Add "Apply" button for code blocks in assistant messages
            if (sub.type == SubBlock::Code && m_applyStorage && !sub.content.trimmed().isEmpty()) {
                // Try to extract file path from preceding prose
                static const QRegularExpression filePathRe(
                    R"((?:`|\b)(/[^\s:`]+\.\w+)(?:`|\b))");
                QString filePath;
                auto fileMatch = filePathRe.match(lastProse);
                if (fileMatch.hasMatch()) {
                    filePath = fileMatch.captured(1);
                }

                int idx = m_applyStorage->size();
                ApplyBlock ab;
                ab.filePath = filePath;
                ab.code = sub.content;
                m_applyStorage->append(ab);

                QString applyLabel = filePath.isEmpty()
                    ? QStringLiteral("✎ Apply")
                    : QStringLiteral("✎ Apply to ") + escapeHtml(filePath.section('/', -1));
                inner += QStringLiteral(
                    "<div style=\"margin:-4px 0 6px 0;\">"
                    "<a href=\"ollie-apply:///") + QString::number(idx) +
                    QStringLiteral("\" style=\"color:#4caf50;text-decoration:none;"
                    "font-size:small;padding:2px 8px;"
                    "border:1px solid #4caf50;border-radius:3px;\">") +
                    applyLabel + QStringLiteral("</a></div>");
            }
        }

        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:8px 0;\">"
            "<tr><td style=\"background:") + assistantBg +
            QStringLiteral(";padding:10px 14px;border:1px solid ") + borderColor +
            QStringLiteral(";\">") + inner +
            QStringLiteral("</td><td width=\"5%\"></td></tr>"
            "<tr><td><span style=\"font-size:small;color:gray;\">Ollie</span></td><td></td></tr>"
            "</table>");
    }
    case ChatBlock::Reasoning: {
        // Show abbreviated with dimmed styling
        QString abbreviated = block.content.left(200);
        if (block.content.size() > 200)
            abbreviated += QStringLiteral("...");

        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:4px 0;\">"
            "<tr><td style=\"background:") + metaBg +
            QStringLiteral(";padding:6px 10px;border-left:3px solid ") + metaColor +
            QStringLiteral(";\">"
            "<div style=\"font-size:small;color:") + metaColor +
            QStringLiteral(";\">▸ reasoning</div>"
            "<pre style=\"margin:4px 0;white-space:pre-wrap;color:") + metaColor +
            QStringLiteral(";font-size:small;\">") + escapeHtml(abbreviated) +
            QStringLiteral("</pre></td></tr></table>");
    }
    case ChatBlock::Call: {
        QString label = QStringLiteral("⚡ call");
        if (!block.name.isEmpty())
            label += QStringLiteral(": ") + escapeHtml(block.name);

        // Render full call args with sub-block parsing (preserve whitespace)
        QList<SubBlock> subs = parseSubBlocks(block.content);
        QString inner;
        for (const auto &sub : subs)
            inner += subBlockToHtml(sub, true);

        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:4px 0;\">"
            "<tr><td style=\"background:") + metaBg +
            QStringLiteral(";padding:6px 10px;border-left:3px solid #ff9800;\">"
            "<div style=\"font-size:small;color:#ff9800;\">") + label +
            QStringLiteral("</div>") + inner +
            QStringLiteral("</td></tr></table>");
    }
    case ChatBlock::Tool: {
        QString label = QStringLiteral("◆ tool");
        if (!block.name.isEmpty())
            label += QStringLiteral(": ") + escapeHtml(block.name);

        // Render full tool output with sub-block parsing (preserve whitespace)
        QList<SubBlock> subs = parseSubBlocks(block.content);
        QString inner;
        for (const auto &sub : subs)
            inner += subBlockToHtml(sub, true);

        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:4px 0;\">"
            "<tr><td style=\"background:") + metaBg +
            QStringLiteral(";padding:6px 10px;border-left:3px solid ") + metaColor +
            QStringLiteral(";\">"
            "<div style=\"font-size:small;color:") + metaColor +
            QStringLiteral(";\">") + label +
            QStringLiteral("</div>") + inner +
            QStringLiteral("</td></tr></table>");
    }
    case ChatBlock::Error: {
        return QStringLiteral(
            "<table width=\"100%\" cellpadding=\"0\" cellspacing=\"0\" style=\"margin:6px 0;\">"
            "<tr><td style=\"background:") + errorBg +
            QStringLiteral(";padding:8px 12px;border-left:3px solid #e53935;\">"
            "<pre style=\"margin:0;white-space:pre-wrap;color:#e53935;\">") +
            escapeHtml(block.content) +
            QStringLiteral("</pre></td></tr></table>");
    }
    case ChatBlock::Info: {
        return QStringLiteral(
            "<div style=\"margin:4px 0;font-style:italic;color:") + metaColor +
            QStringLiteral(";font-size:small;\">") +
            escapeHtml(block.content) + QStringLiteral("</div>");
    }
    }
    return {};
}

// --- Full Render ---

void ChatRenderer::render(QTextEdit *view, const QString &rawText)
{
    detectTheme();

    // Clear storage before re-render (blocks get re-indexed)
    if (m_diffStorage)
        m_diffStorage->clear();
    if (m_applyStorage)
        m_applyStorage->clear();

    QList<ChatBlock> blocks = parseBlocks(rawText);
    QString html = QStringLiteral(
        "<html><body style=\"font-family:sans-serif;font-size:medium;\">");
    for (const auto &block : blocks)
        html += blockToHtml(block);
    html += QStringLiteral("</body></html>");

    auto *sb = view->verticalScrollBar();
    bool atBottom = (sb->maximum() <= 0) || (sb->value() >= sb->maximum() - 4);
    int savedPos = sb->value();

    // Block signals to prevent intermediate scroll events during setHtml
    sb->blockSignals(true);
    view->setHtml(html);

    if (atBottom) {
        // Move cursor to end and ensure visible — this forces layout and
        // scrolls reliably, unlike sb->setValue(sb->maximum()) which can
        // read a stale maximum before relayout completes.
        QTextCursor cursor = view->textCursor();
        cursor.movePosition(QTextCursor::End);
        view->setTextCursor(cursor);
        view->ensureCursorVisible();
    } else {
        // Restore exact scroll position — clamp to new maximum
        sb->setValue(qMin(savedPos, sb->maximum()));
    }
    sb->blockSignals(false);

    m_lastRawText = rawText;
    m_lastBlockCount = blocks.size();
}

bool ChatRenderer::renderIncremental(QTextEdit *view, const QString &rawText)
{
    // If text is just appended, try incremental
    if (!rawText.startsWith(m_lastRawText) || m_lastRawText.isEmpty()) {
        render(view, rawText);
        return false;
    }

    // Full re-render (the HTML approach doesn't support true incremental well
    // since setHtml replaces everything). For performance with typical chat
    // sizes this is fine — QTextEdit handles it efficiently.
    render(view, rawText);
    return true;
}