#include "chatblockmodel.h"

#include <QTextDocument>
#include <QApplication>
#include <QPalette>
#include <QNetworkAccessManager>
#include <QNetworkReply>
#include <QNetworkRequest>
#include <QUrl>
#include <QUrlQuery>
#include <QCryptographicHash>
#include <QJsonDocument>
#include <QJsonObject>
#include <QRegularExpression>

#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 ---

ChatBlock::Type ChatBlockModel::blockTypeFromString(const QString &role)
{
    if (role == "user")           return ChatBlock::User;
    if (role == "assistant")      return ChatBlock::Assistant;
    if (role == "context")        return ChatBlock::Context;
    if (role == "reasoning")      return ChatBlock::Reasoning;
    if (role == "call")           return ChatBlock::Call;
    if (role == "tool")           return ChatBlock::Tool;
    if (role == "error")          return ChatBlock::Error;
    if (role == "info")           return ChatBlock::Info;
    if (role == "retry")          return ChatBlock::Retry;
    if (role == "interrupt")      return ChatBlock::Interrupt;
    return ChatBlock::Default;
}

// Code fence regex (used for content rendering, not block parsing)
static const QRegularExpression s_fenceOpenRe(R"(^```([A-Za-z0-9_+.-]*).*$)");
static const QRegularExpression s_fenceCloseRe(R"(^```$)");

ChatBlockModel::ChatBlockModel(QObject *parent)
    : QAbstractListModel(parent)
    , m_repo(new KSyntaxHighlighting::Repository())
    , m_netManager(new QNetworkAccessManager(this))
{
    connect(m_netManager, &QNetworkAccessManager::finished,
            this, &ChatBlockModel::onDiagramReply);
}

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::Context:   return QStringLiteral("context");
        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 ContentRole:
        return b.content;
    case RenderedHtmlRole:
        return b.renderedHtml;
    case SegmentsRole:
        return segmentsToVariant(b.segments);
    case BlockIdRole:
        return b.blockId;
    case PartialRole:
        return b.partial;
    }
    return {};
}

QVariantList ChatBlockModel::segmentsToVariant(const QList<BlockSegment> &segments) const
{
    QVariantList result;
    for (const BlockSegment &seg : segments) {
        QVariantMap m;
        m[QStringLiteral("type")] = (seg.type == BlockSegment::Text)
            ? QStringLiteral("text") : QStringLiteral("diagram");
        m[QStringLiteral("html")] = seg.html;
        if (seg.type == BlockSegment::Diagram) {
            m[QStringLiteral("lang")] = seg.diagramLang;
            m[QStringLiteral("source")] = seg.diagramSource;
            m[QStringLiteral("svg")] = seg.diagramSvg;
            m[QStringLiteral("error")] = seg.diagramError;
            m[QStringLiteral("loading")] = seg.diagramLoading;
        }
        result.append(m);
    }
    return result;
}

QHash<int, QByteArray> ChatBlockModel::roleNames() const
{
    return {
        {BlockTypeRole, "blockType"},
        {ContentRole,  "blockContent"},
        {RenderedHtmlRole, "blockRenderedHtml"},
        {SegmentsRole, "blockSegments"},
        {BlockIdRole, "blockId"},
        {PartialRole, "blockPartial"},
    };
}

void ChatBlockModel::appendData(const QString &text)
{
    m_lineBuffer += text.toUtf8();

    // Process complete JSONL lines
    int idx;
    while ((idx = m_lineBuffer.indexOf('\n')) >= 0) {
        QByteArray line = m_lineBuffer.left(idx);
        m_lineBuffer = m_lineBuffer.mid(idx + 1);
        if (!line.isEmpty()) {
            processJsonLine(line);
        }
    }
}

void ChatBlockModel::processJsonLine(const QByteArray &line)
{
    QJsonParseError error;
    QJsonDocument doc = QJsonDocument::fromJson(line, &error);
    if (error.error != QJsonParseError::NoError || !doc.isObject()) {
        qWarning() << "ChatBlockModel: invalid JSON line:" << error.errorString();
        return;
    }

    QJsonObject obj = doc.object();
    QString role = obj["role"].toString();
    QString id = obj["id"].toString();
    QString content = obj["content"].toString();
    QString name = obj["name"].toString();
    QString format = obj["format"].toString();
    bool partial = obj["partial"].toBool(false);

    // Skip context blocks from rendering (context is for LLM, not display)
    if (role == "context") {
        return;
    }

    // Find existing block with same ID (for streaming updates)
    int existingIdx = indexForBlockId(id);

    if (existingIdx >= 0) {
        // Update existing block
        ChatBlock &block = m_blocks[existingIdx];
        block.content = content;
        block.partial = partial;
        rebuildBlockHtml(existingIdx);

        QModelIndex idx = index(existingIdx);
        emit dataChanged(idx, idx, {ContentRole, RenderedHtmlRole, SegmentsRole, PartialRole});
    } else {
        // New block
        int newIdx = m_blocks.size();
        beginInsertRows(QModelIndex(), newIdx, newIdx);

        ChatBlock block;
        block.type = blockTypeFromString(role);
        block.blockId = id;
        block.content = content;
        block.name = name;
        block.format = format;
        block.partial = partial;

        m_blocks.append(block);
        endInsertRows();
        emit countChanged();

        // Render after insertion so rebuildBlockHtml can find it
        rebuildBlockHtml(newIdx);
        QModelIndex idx = index(newIdx);
        emit dataChanged(idx, idx, {ContentRole, RenderedHtmlRole, SegmentsRole});
    }
}

void ChatBlockModel::reset()
{
    beginResetModel();
    m_blocks.clear();
    m_lineBuffer.clear();
    m_pendingDiagrams.clear();
    m_inFlightRequests.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();
}

int ChatBlockModel::indexForBlockId(const QString &blockId) const
{
    for (int i = 0; i < m_blocks.size(); ++i) {
        if (m_blocks[i].blockId == blockId)
            return i;
    }
    return -1;
}

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::highlightCodeWrapped(const QString &innerHtml) const
{
    return QStringLiteral("<pre style=\"white-space:pre-wrap;\">") + innerHtml + QStringLiteral("</pre>");
}

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

bool ChatBlockModel::isTableLine(const QString &line)
{
    const QString trimmed = line.trimmed();
    return trimmed.startsWith('|') && trimmed.endsWith('|') && trimmed.count('|') >= 2;
}

QString ChatBlockModel::renderTable(const QStringList &tableLines) const
{
    const QString defaultColor = getDefaultTextColor();

    QString html = QStringLiteral("<table border=\"1\" cellspacing=\"0\" cellpadding=\"3\" style=\"color:")
        + defaultColor + QStringLiteral(";\">");

    bool pastSeparator = false;

    for (const QString &line : tableLines) {
        const QString trimmed = line.trimmed();
        // Skip separator rows (|---|---|)
        static QRegularExpression sepRe(QStringLiteral(R"(^\|[\s:|-]+\|$)"));
        if (sepRe.match(trimmed).hasMatch()) {
            pastSeparator = true;
            continue;
        }

        // Strip leading/trailing |, split by |
        QString inner = trimmed;
        if (inner.startsWith('|')) inner = inner.mid(1);
        if (inner.endsWith('|')) inner.chop(1);
        const QStringList cells = inner.split('|');

        const QString tag = pastSeparator ? QStringLiteral("td") : QStringLiteral("th");

        html += QStringLiteral("<tr>");
        for (const QString &cell : cells) {
            html += QStringLiteral("<") + tag + QStringLiteral(">")
                + cell.trimmed().toHtmlEscaped()
                + QStringLiteral("</") + tag + QStringLiteral(">");
        }
        html += QStringLiteral("</tr>");
    }

    html += QStringLiteral("</table>");
    return html;
}

QString ChatBlockModel::renderProse(const QStringList &lines) const
{
    QString html;
    QStringList textRun;
    QStringList tableRun;

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

    auto flushTable = [&]() {
        if (!tableRun.isEmpty()) {
            html += renderTable(tableRun);
            tableRun.clear();
        }
    };

    for (const QString &line : lines) {
        if (isTableLine(line)) {
            flushText();
            tableRun.append(line);
        } else {
            flushTable();
            textRun.append(line);
        }
    }

    flushText();
    flushTable();
    return html;
}

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

    ChatBlock &block = m_blocks[blockIndex];

    // For call blocks, show tool name and JSON formatted args
    if (block.type == ChatBlock::Call) {
        QString html = QStringLiteral("<div style=\"color:") + getDefaultTextColor() + QStringLiteral(";\">");
        if (!block.name.isEmpty()) {
            html += QStringLiteral("<b>") + block.name.toHtmlEscaped() + QStringLiteral("</b><br/>");
        }
        html += highlightCode(block.content, QStringLiteral("json"));
        html += QStringLiteral("</div>");
        block.renderedHtml = html;

        BlockSegment seg;
        seg.type = BlockSegment::Text;
        seg.html = html;
        block.segments = {seg};
        return;
    }

    // For tool blocks, highlight with format hint
    if (block.type == ChatBlock::Tool) {
        QString lang = block.format.isEmpty() ? QStringLiteral("text") : block.format;
        block.renderedHtml = highlightCode(block.content, lang);

        BlockSegment seg;
        seg.type = BlockSegment::Text;
        seg.html = block.renderedHtml;
        block.segments = {seg};
        return;
    }

    // For other blocks, render content with code fence detection
    QString html;
    QList<BlockSegment> segments;

    // 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()) {
            QString proseHtml = renderProse(proseLines);
            html += proseHtml;

            BlockSegment seg;
            seg.type = BlockSegment::Text;
            seg.html = proseHtml;
            segments.append(seg);

            proseLines.clear();
        }
    };

    auto flushFence = [&](bool complete) {
        if (fenceLines.isEmpty()) {
            fenceLang.clear();
            return;
        }

        const QString code = fenceLines.join('\n');

        if (isDiagramLanguage(fenceLang)) {
            BlockSegment seg;
            seg.type = BlockSegment::Diagram;
            seg.diagramLang = fenceLang;
            seg.diagramSource = code;
            seg.html = highlightCode(code, fenceLang);  // Source view HTML

            if (!complete) {
                // Streaming — show placeholder, don't request yet
                html += diagramPlaceholder(fenceLang);
                seg.diagramLoading = true;
            } else {
                // Complete fence — render through the local diagram server.
                html += diagramPlaceholder(fenceLang);
                seg.diagramLoading = true;
                requestDiagram(fenceLang, code, blockIndex);
            }
            segments.append(seg);
        } else {
            QString codeHtml = highlightCode(code, fenceLang);
            html += codeHtml;

            BlockSegment seg;
            seg.type = BlockSegment::Text;
            seg.html = codeHtml;
            segments.append(seg);
        }

        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(true);  // complete
                inFence = false;
                continue;
            }
            fenceLines.append(line);
        }
    }

    // Handle any remaining content
    if (inFence) {
        // Unclosed fence (streaming)
        flushFence(false);
    } else {
        flushProse();
    }

    block.renderedHtml = html;
    block.segments = segments;
}

void ChatBlockModel::renderLastSegmentPartial(int blockIndex, const QString &partialLine)
{
    if (blockIndex < 0 || blockIndex >= m_blocks.size())
        return;

    ChatBlock &block = m_blocks[blockIndex];
    
    // Append partial line to content and re-render
    QString tempContent = block.content;
    if (!tempContent.isEmpty() && !tempContent.endsWith('\n')) {
        tempContent += partialLine;
    } else {
        tempContent += partialLine;
    }
    
    block.content = tempContent;
    rebuildBlockHtml(blockIndex);
}

void ChatBlockModel::rebuildRenderedHtmlFromSegments(ChatBlock &block)
{
    QString html;
    for (const BlockSegment &seg : block.segments) {
        if (seg.type == BlockSegment::Diagram) {
            if (!seg.diagramSvg.isEmpty()) {
                html += seg.diagramSvg;
            } else if (!seg.diagramError.isEmpty()) {
                html += QStringLiteral("<div style=\"color:red;\">") + seg.diagramError.toHtmlEscaped() + QStringLiteral("</div>");
            } else {
                html += QStringLiteral("<div>[Loading diagram...]</div>");
            }
        } else {
            html += seg.html;
        }
    }
    block.renderedHtml = html;
}

// --- Diagram support ---

bool ChatBlockModel::isDiagramLanguage(const QString &lang) const
{
    return lang == QStringLiteral("plantuml") ||
           lang == QStringLiteral("mermaid") ||
           lang == QStringLiteral("dot") ||
           lang == QStringLiteral("graphviz");
}

QString ChatBlockModel::diagramPlaceholder(const QString &lang) const
{
    return QStringLiteral("<div style=\"padding:8px;background:#f0f0f0;border:1px solid #ccc;\">")
        + QStringLiteral("[Rendering ") + lang.toHtmlEscaped() + QStringLiteral(" diagram...]</div>");
}

QString ChatBlockModel::diagramErrorHtml(const QString &error) const
{
    return QStringLiteral("<div style=\"padding:8px;background:#fee;border:1px solid #c00;color:#c00;\">")
        + QStringLiteral("Diagram error: ") + error.toHtmlEscaped() + QStringLiteral("</div>");
}

QString ChatBlockModel::diagramSvgHtml(const QByteArray &svg) const
{
    return QStringLiteral("<img src=\"data:image/svg+xml;base64,")
        + QString::fromLatin1(svg.toBase64())
        + QStringLiteral("\" style=\"max-width:100%;\" />");
}

QString ChatBlockModel::diagramPngHtml(const QByteArray &png) const
{
    return QStringLiteral("<img src=\"data:image/png;base64,")
        + QString::fromLatin1(png.toBase64())
        + QStringLiteral("\" style=\"max-width:100%;\" />");
}

QString ChatBlockModel::diagramJpegHtml(const QByteArray &jpeg) const
{
    return QStringLiteral("<img src=\"data:image/jpeg;base64,")
        + QString::fromLatin1(jpeg.toBase64())
        + QStringLiteral("\" style=\"max-width:100%;\" />");
}

void ChatBlockModel::requestDiagram(const QString &lang, const QString &code, int blockIndex)
{
    // Hash the request for deduplication
    const QString hashInput = lang + QStringLiteral(":") + code;
    const QString hash = QString::fromLatin1(
        QCryptographicHash::hash(hashInput.toUtf8(), QCryptographicHash::Sha256).toHex().left(16));

    // Track which blocks want this diagram
    m_pendingDiagrams[hash].append(blockIndex);

    // Don't re-request if already in flight
    if (m_inFlightRequests.contains(hash))
        return;
    m_inFlightRequests.insert(hash);

    // Request from local diagram renderer (ollie-diagram service)
    QUrl url(QStringLiteral("http://localhost:18019/render"));
    QUrlQuery query;
    query.addQueryItem(QStringLiteral("lang"), lang);
    url.setQuery(query);

    QNetworkRequest request(url);
    request.setHeader(QNetworkRequest::ContentTypeHeader, QStringLiteral("text/plain"));
    request.setAttribute(QNetworkRequest::User, hash);

    m_netManager->post(request, code.toUtf8());
}

void ChatBlockModel::onDiagramReply(QNetworkReply *reply)
{
    reply->deleteLater();

    const QString hash = reply->request().attribute(QNetworkRequest::User).toString();
    m_inFlightRequests.remove(hash);

    QList<int> indices = m_pendingDiagrams.take(hash);
    if (indices.isEmpty())
        return;

    QString resultHtml;
    if (reply->error() != QNetworkReply::NoError) {
        resultHtml = diagramErrorHtml(reply->errorString());
    } else {
        const QByteArray data = reply->readAll();
        const QString contentType = reply->header(QNetworkRequest::ContentTypeHeader).toString();

        if (contentType.contains(QStringLiteral("svg"))) {
            resultHtml = diagramSvgHtml(data);
        } else if (contentType.contains(QStringLiteral("png"))) {
            resultHtml = diagramPngHtml(data);
        } else if (contentType.contains(QStringLiteral("jpeg")) || contentType.contains(QStringLiteral("jpg"))) {
            resultHtml = diagramJpegHtml(data);
        } else {
            resultHtml = diagramErrorHtml(QStringLiteral("Unknown content type: ") + contentType);
        }
    }

    // Update all blocks waiting for this diagram
    for (int idx : indices) {
        if (idx < 0 || idx >= m_blocks.size())
            continue;

        ChatBlock &block = m_blocks[idx];
        for (BlockSegment &seg : block.segments) {
            if (seg.type == BlockSegment::Diagram && seg.diagramLoading) {
                seg.diagramSvg = resultHtml;
                seg.diagramLoading = false;
                if (reply->error() != QNetworkReply::NoError) {
                    seg.diagramError = reply->errorString();
                }
            }
        }
        rebuildRenderedHtmlFromSegments(block);

        QModelIndex modelIdx = index(idx);
        emit dataChanged(modelIdx, modelIdx, {RenderedHtmlRole, SegmentsRole});
    }
}