#include "chatblockmodel.h"

#include <QTextDocument>
#include <QRegularExpressionMatch>
#include <QApplication>
#include <QPalette>
#include <QNetworkAccessManager>
#include <QNetworkReply>
#include <QNetworkRequest>
#include <QUrl>
#include <QUrlQuery>
#include <QCryptographicHash>

#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())
    , 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::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;
    case SegmentsRole:
        return segmentsToVariant(b.segments);
    }
    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"},
        {BlockNameRole, "blockName"},
        {ContentRole,  "blockContent"},
        {RenderedHtmlRole, "blockRenderedHtml"},
        {SegmentsRole, "blockSegments"},
    };
}

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();
    m_pendingDiagrams.clear();
    m_inFlightRequests.clear();
    m_diagramCache.clear();  // Clear cache on reset to pick up fixes
    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();
}

void ChatBlockModel::clearDiagramCache()
{
    m_diagramCache.clear();
    qDebug() << "Diagram cache cleared";
}

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;
    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 = escapeForHtml(proseLines.join('\n'));
            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 — check cache or request diagram
                const QString cacheKey = fenceLang + QStringLiteral(":") + 
                    QString::fromLatin1(QCryptographicHash::hash(
                        code.toUtf8(), QCryptographicHash::Md5).toHex());
                
                if (m_diagramCache.contains(cacheKey)) {
                    const QByteArray cached = m_diagramCache.value(cacheKey);
                    if (cached.startsWith("ERROR:")) {
                        // Error: fall back to source code rendering (not a diagram segment)
                        QString codeHtml = highlightCode(code, fenceLang);
                        html += codeHtml;
                        
                        BlockSegment textSeg;
                        textSeg.type = BlockSegment::Text;
                        textSeg.html = codeHtml;
                        segments.append(textSeg);
                        
                        fenceLines.clear();
                        fenceLang.clear();
                        return;  // Don't append diagram segment
                    } else {
                        // Check if JPEG, PNG, or SVG
                        if (cached.startsWith("JPEG:")) {
                            QByteArray jpegData = cached.mid(5);  // Skip "JPEG:" prefix
                            html += diagramJpegHtml(jpegData);
                            seg.diagramSvg = QStringLiteral("data:image/jpeg;base64,") + 
                                QString::fromLatin1(jpegData.toBase64());
                        } else if (cached.startsWith("PNG:")) {
                            QByteArray pngData = cached.mid(4);  // Skip "PNG:" prefix
                            html += diagramPngHtml(pngData);
                            seg.diagramSvg = QStringLiteral("data:image/png;base64,") + 
                                QString::fromLatin1(pngData.toBase64());
                        } else if (cached.startsWith("SVG:")) {
                            QByteArray svgData = cached.mid(4);  // Skip "SVG:" prefix
                            html += diagramSvgHtml(svgData);
                            seg.diagramSvg = QStringLiteral("data:image/svg+xml;base64,") + 
                                QString::fromLatin1(svgData.toBase64());
                        } else {
                            // Legacy: assume SVG if no prefix
                            html += diagramSvgHtml(cached);
                            seg.diagramSvg = QStringLiteral("data:image/svg+xml;base64,") + 
                                QString::fromLatin1(cached.toBase64());
                        }
                    }
                } else {
                    html += diagramPlaceholder(fenceLang);
                    seg.diagramLoading = true;
                    // Request diagram asynchronously
                    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 unclosed fence (streaming) or remaining prose
    if (inFence) {
        // Still inside a fence — render what we have (incomplete)
        flushFence(false);
    } else {
        flushProse();
    }

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

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

// --- Diagram rendering ---

bool ChatBlockModel::isDiagramLanguage(const QString &lang) const
{
    const QString l = lang.toLower();
    return l == QStringLiteral("plantuml") || 
           l == QStringLiteral("mermaid") ||
           l == QStringLiteral("puml");
}

QString ChatBlockModel::diagramPlaceholder(const QString &lang) const
{
    const QString defaultColor = getDefaultTextColor();
    return QStringLiteral("<div style=\"padding:8px;border:1px dashed ") + defaultColor +
           QStringLiteral(";color:") + defaultColor + 
           QStringLiteral(";font-style:italic;\">Loading ") + 
           lang.toHtmlEscaped() + QStringLiteral(" diagram...</div>");
}

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

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

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

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

void ChatBlockModel::requestDiagram(const QString &lang, const QString &code, int blockIndex)
{
    const QString cacheKey = lang + QStringLiteral(":") + 
        QString::fromLatin1(QCryptographicHash::hash(
            code.toUtf8(), QCryptographicHash::Md5).toHex());
    
    // Already have it cached?
    if (m_diagramCache.contains(cacheKey))
        return;
    
    // Track which blocks need this diagram
    if (!m_pendingDiagrams.contains(cacheKey)) {
        m_pendingDiagrams[cacheKey] = QList<int>();
    }
    if (!m_pendingDiagrams[cacheKey].contains(blockIndex)) {
        m_pendingDiagrams[cacheKey].append(blockIndex);
    }
    
    // Already requesting?
    if (m_inFlightRequests.contains(cacheKey))
        return;
    
    m_inFlightRequests.insert(cacheKey);
    
    // Build request
    QNetworkRequest request;
    request.setAttribute(QNetworkRequest::User, cacheKey);
    request.setTransferTimeout(10000);
    
    const QString langLower = lang.toLower();
    
    if (langLower == QStringLiteral("plantuml") || langLower == QStringLiteral("puml")) {
        // PlantUML server accepts POST with raw source
        request.setUrl(QUrl(QStringLiteral("http://localhost:8080/svg")));
        request.setHeader(QNetworkRequest::ContentTypeHeader, QStringLiteral("text/plain"));
        m_netManager->post(request, code.toUtf8());
    } else if (langLower == QStringLiteral("mermaid")) {
        // Mermaid-ink expects base64 in URL path, /img/ for PNG output
        // (SVG uses foreignObject which Qt's SVG renderer doesn't support)
        const QByteArray encoded = code.toUtf8().toBase64();
        request.setUrl(QUrl(QStringLiteral("http://localhost:8081/img/") + QString::fromLatin1(encoded)));
        m_netManager->get(request);
    } else {
        m_inFlightRequests.remove(cacheKey);
        return;
    }
}

void ChatBlockModel::onDiagramReply(QNetworkReply *reply)
{
    reply->deleteLater();
    
    const QString cacheKey = reply->request().attribute(QNetworkRequest::User).toString();
    m_inFlightRequests.remove(cacheKey);
    
    QByteArray imageData;
    bool success = false;
    QString errorMsg;
    
    if (reply->error() == QNetworkReply::NoError) {
        imageData = reply->readAll();
        
        // Check for JPEG/PNG (mermaid) or SVG (plantuml)
        if (imageData.startsWith("\xff\xd8\xff")) {
            // JPEG image
            success = true;
            imageData.prepend("JPEG:");
        } else if (imageData.startsWith("\x89PNG")) {
            // PNG image
            success = true;
            imageData.prepend("PNG:");
        } else if (imageData.contains("<svg") || imageData.startsWith("<?xml")) {
            // SVG image
            success = true;
            imageData.prepend("SVG:");
        } else {
            errorMsg = QStringLiteral("Invalid response (not SVG, PNG, or JPEG)");
        }
    } else {
        errorMsg = reply->errorString();
    }
    
    // Update cache
    if (success) {
        m_diagramCache[cacheKey] = imageData;
    } else {
        // Cache error as empty data with special marker
        m_diagramCache[cacheKey] = QByteArray("ERROR:") + errorMsg.toUtf8();
    }
    
    // Rebuild all blocks waiting for this diagram
    const QList<int> blocks = m_pendingDiagrams.take(cacheKey);
    for (int blockIndex : blocks) {
        if (blockIndex < m_blocks.size()) {
            rebuildBlockHtml(blockIndex);
            QModelIndex idx = index(blockIndex);
            emit dataChanged(idx, idx, {RenderedHtmlRole, SegmentsRole});
        }
    }
}