// ==========================================
// KaTeX Math Renderer — Browser Extension
// ==========================================
// Renders ```latex / ```math code blocks and inline $...$ / $$...$$ syntax.
// Loads KaTeX dynamically on first use.
// ==========================================

Extensions.register({
    id: 'katex-renderer',

    _katexReady: false,
    _katexLoading: null,

    async init(ctx) {
        const self = this;

        // ── Inject styles ──
        this._injectStyles();

        // ── Block renderer: match ```latex or ```math ──
        ctx.renderers.register('katex-block', {
            type: 'block',
            priority: 10,
            match(lang) {
                const l = (lang || '').toLowerCase();
                return l === 'latex' || l === 'math' || l === 'tex' || l === 'katex';
            },
            render(lang, code, container) {
                const id = 'katex-' + Math.random().toString(36).slice(2, 9);
                container.innerHTML = `
                    <div class="katex-block ext-rendered" data-katex-id="${id}">
                        <div class="katex-display-container" data-katex-src="${encodeURIComponent(code.trim())}" data-katex-display="true">
                            <div class="katex-loading">
                                <span class="katex-spinner"></span> Rendering math…
                            </div>
                        </div>
                        <details class="katex-source">
                            <summary>𝑓(𝑥) View source</summary>
                            <pre><code class="language-latex">${self._escapeHtml(code.trim())}</code></pre>
                        </details>
                    </div>
                `;
            }
        });

        // ── Post renderer: render KaTeX blocks + find inline math ──
        ctx.renderers.register('katex-post', {
            type: 'post',
            priority: 10,
            render(container) {
                // Render code blocks with data-katex-src
                const blocks = container.querySelectorAll('.katex-display-container[data-katex-src]');
                if (blocks.length > 0) {
                    blocks.forEach(el => {
                        if (el.dataset.rendered) return;
                        el.dataset.rendered = 'pending';
                        self._renderBlock(el);
                    });
                }

                // Process inline math in text content
                self._processInlineMath(container);
            }
        });

        // Pre-load KaTeX
        this._loadKaTeX();
    },

    async _renderBlock(el) {
        const code = decodeURIComponent(el.dataset.katexSrc);
        const displayMode = el.dataset.katexDisplay === 'true';

        try {
            await this._loadKaTeX();
            el.innerHTML = katex.renderToString(code, {
                displayMode,
                throwOnError: false,
                trust: false,
                strict: false,
            });
            el.dataset.rendered = 'true';
        } catch (e) {
            el.innerHTML = `
                <div class="katex-error">
                    <strong>Math error:</strong> ${this._escapeHtml(e.message || String(e))}
                </div>
            `;
            el.dataset.rendered = 'error';
        }
    },

    async _processInlineMath(container) {
        // Skip if no $ signs present at all
        if (!container.textContent || !container.textContent.includes('$')) return;

        try {
            await this._loadKaTeX();
        } catch {
            return;
        }

        // Walk text nodes, skip code/pre/katex-already-rendered
        const walker = document.createTreeWalker(
            container,
            NodeFilter.SHOW_TEXT,
            {
                acceptNode(node) {
                    const parent = node.parentElement;
                    if (!parent) return NodeFilter.FILTER_REJECT;
                    // Skip code, pre, already-rendered katex, and script elements
                    const tag = parent.tagName;
                    if (tag === 'CODE' || tag === 'PRE' || tag === 'SCRIPT' ||
                        tag === 'TEXTAREA' || tag === 'STYLE') {
                        return NodeFilter.FILTER_REJECT;
                    }
                    if (parent.closest('.katex, .katex-block, .katex-display-container, code, pre')) {
                        return NodeFilter.FILTER_REJECT;
                    }
                    if (!node.textContent.includes('$')) return NodeFilter.FILTER_REJECT;
                    return NodeFilter.FILTER_ACCEPT;
                }
            }
        );

        const textNodes = [];
        let node;
        while (node = walker.nextNode()) textNodes.push(node);

        for (const textNode of textNodes) {
            this._replaceInlineMath(textNode);
        }
    },

    _replaceInlineMath(textNode) {
        const text = textNode.textContent;
        // Match $$...$$ (display) and $...$ (inline), non-greedy
        // Avoid matching escaped \$ or empty $$
        const pattern = /\$\$([^$]+?)\$\$|\$([^$\n]+?)\$/g;
        let match;
        const parts = [];
        let lastIndex = 0;

        while ((match = pattern.exec(text)) !== null) {
            // Text before the match
            if (match.index > lastIndex) {
                parts.push({ type: 'text', value: text.slice(lastIndex, match.index) });
            }

            const displayMode = !!match[1]; // $$ ... $$
            const expr = match[1] || match[2];

            try {
                const html = katex.renderToString(expr.trim(), {
                    displayMode,
                    throwOnError: false,
                    trust: false,
                    strict: false,
                });
                parts.push({ type: 'katex', value: html, displayMode });
            } catch {
                // Leave as-is on error
                parts.push({ type: 'text', value: match[0] });
            }

            lastIndex = match.index + match[0].length;
        }

        if (parts.length === 0) return; // No math found

        // Remaining text after last match
        if (lastIndex < text.length) {
            parts.push({ type: 'text', value: text.slice(lastIndex) });
        }

        // Replace text node with a span containing the rendered parts
        const span = document.createElement('span');
        span.className = 'katex-inline-container';
        for (const part of parts) {
            if (part.type === 'text') {
                span.appendChild(document.createTextNode(part.value));
            } else {
                const wrapper = document.createElement('span');
                wrapper.className = part.displayMode ? 'katex-display' : 'katex-inline';
                wrapper.innerHTML = part.value;
                span.appendChild(wrapper);
            }
        }

        textNode.parentNode.replaceChild(span, textNode);
    },

    _loadKaTeX() {
        if (this._katexReady) return Promise.resolve();
        if (this._katexLoading) return this._katexLoading;

        this._katexLoading = new Promise((resolve, reject) => {
            if (typeof katex !== 'undefined') {
                this._katexReady = true;
                resolve();
                return;
            }

            const base = (window.__BASE__ || '');
            const localCSS = `${base}/vendor/katex/katex.min.css`;
            const localJS = `${base}/vendor/katex/katex.min.js`;
            const cdnCSS = 'https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css';
            const cdnJS = 'https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.js';

            // Load CSS first
            const link = document.createElement('link');
            link.rel = 'stylesheet';
            link.href = localCSS;
            link.onerror = () => {
                link.href = cdnCSS;
            };
            document.head.appendChild(link);

            // Load JS
            const script = document.createElement('script');
            script.src = localJS;
            script.onload = () => {
                this._katexReady = true;
                resolve();
            };
            script.onerror = () => {
                console.warn('[KaTeX] Local vendor not found, trying CDN');
                const cdn = document.createElement('script');
                cdn.src = cdnJS;
                cdn.onload = () => {
                    this._katexReady = true;
                    resolve();
                };
                cdn.onerror = () => {
                    console.error('[KaTeX] Failed to load from both local and CDN');
                    reject(new Error('Failed to load KaTeX'));
                };
                document.head.appendChild(cdn);
            };
            document.head.appendChild(script);
        });

        return this._katexLoading;
    },

    _escapeHtml(str) {
        const div = document.createElement('div');
        div.textContent = str;
        return div.innerHTML;
    },

    _injectStyles() {
        if (document.getElementById('ext-style-katex-renderer')) return;
        const style = document.createElement('style');
        style.id = 'ext-style-katex-renderer';
        style.textContent = `
.katex-block {
    background: var(--bg-2); border: 1px solid var(--border);
    border-radius: 8px; overflow: hidden; margin: 12px 0;
}
.katex-display-container {
    padding: 16px; text-align: center; min-height: 40px;
}
.katex-loading {
    color: var(--text-3); font-size: 13px; padding: 20px;
    display: flex; align-items: center; justify-content: center; gap: 8px;
}
.katex-spinner {
    display: inline-block; width: 14px; height: 14px;
    border: 2px solid var(--border); border-top-color: var(--accent);
    border-radius: 50%; animation: katex-spin 0.8s linear infinite;
}
@keyframes katex-spin { to { transform: rotate(360deg); } }
.katex-error {
    color: var(--error, #e74c3c); background: var(--bg-3, rgba(231,76,60,0.1));
    padding: 12px 16px; border-radius: 4px; font-size: 13px;
    font-family: var(--mono); white-space: pre-wrap;
}
.katex-source { border-top: 1px solid var(--border); font-size: 12px; }
.katex-source summary {
    padding: 6px 12px; cursor: pointer; color: var(--text-3); user-select: none;
}
.katex-source summary:hover { color: var(--text-2); }
.katex-source pre {
    margin: 0; border-radius: 0; border: none;
    max-height: 200px; overflow: auto;
}
.katex-inline-container { display: inline; }
.katex-inline .katex { font-size: 1.05em; }
.katex-display { display: block; text-align: center; margin: 8px 0; }`;
        document.head.appendChild(style);
    },

    destroy() {
        document.getElementById('ext-style-katex-renderer')?.remove();
    }
});
