const { describe, it } = require('node:test');
const assert = require('node:assert');
const fs = require('fs');
const path = require('path');
const { createBrowserContext, loadSource, SRC } = require('./helpers');
const vm = require('vm');

// Skip all tests if the source file doesn't exist.
const EXTENSIONS_EXISTS = fs.existsSync(path.join(SRC, 'extensions.js'));
const maybeDescribe = EXTENSIONS_EXISTS ? describe : describe.skip;

/**
 * Load Events + Extensions into a browser context.
 */
function loadExtensions(overrides = {}) {
    const sandbox = createBrowserContext(overrides);
    loadSource(sandbox, 'sb.js');
    loadSource(sandbox, 'events.js');
    loadSource(sandbox, 'extensions.js');
    const ctx = vm.createContext(sandbox);
    return {
        Extensions: vm.runInContext('Extensions', ctx),
        Events: vm.runInContext('Events', ctx),
        sandbox,
    };
}

// ═══════════════════════════════════════════
// KaTeX Renderer Extension
// ═══════════════════════════════════════════

maybeDescribe('KaTeX renderer extension', () => {
    function loadWithKaTeX() {
        const ctx = loadExtensions();
        ctx.Extensions.register({
            id: 'katex-renderer',
            _katexReady: false,
            _katexLoading: null,
            _escapeHtml(str) {
                return str.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;');
            },
            init(extCtx) {
                const self = this;
                extCtx.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) {
                        container.innerHTML = `
                            <div class="katex-block ext-rendered">
                                <div class="katex-display-container" data-katex-src="${encodeURIComponent(code.trim())}" data-katex-display="true">
                                    <div class="katex-loading">Rendering math…</div>
                                </div>
                                <details class="katex-source">
                                    <summary>View source</summary>
                                    <pre><code>${self._escapeHtml(code.trim())}</code></pre>
                                </details>
                            </div>
                        `;
                    }
                });
                extCtx.renderers.register('katex-post', {
                    type: 'post',
                    priority: 10,
                    render(container) {
                        container._katexPostRan = true;
                    }
                });
            },
            destroy() {}
        });
        return ctx;
    }

    it('registers both block and post renderers', async () => {
        const ctx = loadWithKaTeX();
        await ctx.Extensions.initAll();
        const info = ctx.Extensions.debug();
        assert.strictEqual(info.extensions['katex-renderer'].active, true);
        assert.strictEqual(info.extensions['katex-renderer'].renderers.length, 2);
    });

    it('matches latex, math, tex, and katex languages', async () => {
        const ctx = loadWithKaTeX();
        await ctx.Extensions.initAll();

        for (const lang of ['latex', 'math', 'tex', 'katex']) {
            const container = { innerHTML: '' };
            const handled = ctx.Extensions.runBlockRenderers(lang, 'E = mc^2', container);
            assert.ok(handled, `should match language: ${lang}`);
            assert.ok(container.innerHTML.includes('katex-block'), `should render katex block for: ${lang}`);
        }
    });

    it('does not match other languages', async () => {
        const ctx = loadWithKaTeX();
        await ctx.Extensions.initAll();

        for (const lang of ['python', 'javascript', 'mermaid', 'csv']) {
            const container = { innerHTML: '' };
            const handled = ctx.Extensions.runBlockRenderers(lang, 'code', container);
            assert.strictEqual(handled, false, `should not match language: ${lang}`);
        }
    });

    it('includes encoded source in data attribute', async () => {
        const ctx = loadWithKaTeX();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('latex', '\\frac{a}{b}', container);
        assert.ok(container.innerHTML.includes(encodeURIComponent('\\frac{a}{b}')));
    });

    it('escapes HTML in source view', async () => {
        const ctx = loadWithKaTeX();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('math', '<script>alert(1)</script>', container);
        assert.ok(container.innerHTML.includes('&lt;script&gt;'));
        assert.ok(!container.innerHTML.includes('<script>alert'));
    });

    it('post renderer runs on container', async () => {
        const ctx = loadWithKaTeX();
        await ctx.Extensions.initAll();

        const container = {};
        ctx.Extensions.runPostRenderers(container);
        assert.strictEqual(container._katexPostRan, true);
    });
});

// ═══════════════════════════════════════════
// CSV Table Extension
// ═══════════════════════════════════════════

maybeDescribe('CSV Table extension', () => {
    function loadWithCSV() {
        const ctx = loadExtensions();
        ctx.Extensions.register({
            id: 'csv-table',
            _escapeHtml(str) {
                return str.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;');
            },
            _parseCSV(text, delimiter) {
                const rows = [];
                let row = [];
                let field = '';
                let inQuotes = false;
                let i = 0;
                while (i < text.length) {
                    const ch = text[i];
                    if (inQuotes) {
                        if (ch === '"') {
                            if (i + 1 < text.length && text[i + 1] === '"') {
                                field += '"'; i += 2;
                            } else {
                                inQuotes = false; i++;
                            }
                        } else { field += ch; i++; }
                    } else {
                        if (ch === '"' && field === '') { inQuotes = true; i++; }
                        else if (ch === delimiter) { row.push(field.trim()); field = ''; i++; }
                        else if (ch === '\n' || (ch === '\r' && text[i + 1] === '\n')) {
                            row.push(field.trim());
                            if (row.some(c => c !== '')) rows.push(row);
                            row = []; field = '';
                            i += (ch === '\r') ? 2 : 1;
                        } else { field += ch; i++; }
                    }
                }
                row.push(field.trim());
                if (row.some(c => c !== '')) rows.push(row);
                if (rows.length > 0) {
                    const maxCols = Math.max(...rows.map(r => r.length));
                    rows.forEach(r => { while (r.length < maxCols) r.push(''); });
                }
                return rows;
            },
            init(extCtx) {
                const self = this;
                extCtx.renderers.register('csv-block', {
                    type: 'block',
                    priority: 10,
                    match(lang) {
                        const l = (lang || '').toLowerCase();
                        return l === 'csv' || l === 'tsv';
                    },
                    render(lang, code, container) {
                        const delimiter = lang.toLowerCase() === 'tsv' ? '\t' : ',';
                        const rows = self._parseCSV(code.trim(), delimiter);
                        if (rows.length === 0) {
                            container.innerHTML = '<div class="csv-empty">No data</div>';
                            return;
                        }
                        const headers = rows[0];
                        const data = rows.slice(1);
                        container.innerHTML = `
                            <div class="csv-table-block ext-rendered">
                                <div class="csv-toolbar">
                                    <span class="csv-info">${data.length} rows × ${headers.length} cols</span>
                                </div>
                                <table class="csv-table">
                                    <thead><tr>${headers.map((h, i) => `<th data-col="${i}">${self._escapeHtml(h)}</th>`).join('')}</tr></thead>
                                    <tbody>${data.map(row => `<tr>${row.map(c => `<td>${self._escapeHtml(c)}</td>`).join('')}</tr>`).join('')}</tbody>
                                </table>
                            </div>
                        `;
                    }
                });
            },
            destroy() {}
        });
        return ctx;
    }

    it('matches csv and tsv languages', async () => {
        const ctx = loadWithCSV();
        await ctx.Extensions.initAll();

        for (const lang of ['csv', 'tsv']) {
            const container = { innerHTML: '' };
            const handled = ctx.Extensions.runBlockRenderers(lang, 'a,b\n1,2', container);
            assert.ok(handled, `should match language: ${lang}`);
        }
    });

    it('does not match other languages', async () => {
        const ctx = loadWithCSV();
        await ctx.Extensions.initAll();

        const handled = ctx.Extensions.runBlockRenderers('json', '{}', { innerHTML: '' });
        assert.strictEqual(handled, false);
    });

    it('renders headers and data rows', async () => {
        const ctx = loadWithCSV();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('csv', 'Name,Age\nAlice,30\nBob,25', container);
        assert.ok(container.innerHTML.includes('csv-table'));
        assert.ok(container.innerHTML.includes('Name'));
        assert.ok(container.innerHTML.includes('Alice'));
        assert.ok(container.innerHTML.includes('30'));
        assert.ok(container.innerHTML.includes('2 rows'));
    });

    it('handles quoted CSV fields', async () => {
        const ctx = loadWithCSV();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('csv', 'Name,Note\n"Smith, John","Has a ""quote"""', container);
        assert.ok(container.innerHTML.includes('Smith, John'));
        assert.ok(container.innerHTML.includes('Has a "quote"'));
    });

    it('handles empty input', async () => {
        const ctx = loadWithCSV();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('csv', '', container);
        assert.ok(container.innerHTML.includes('csv-empty') || container.innerHTML.includes('No data'));
    });

    it('escapes HTML in cell content', async () => {
        const ctx = loadWithCSV();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('csv', 'H1\n<img onerror=alert(1)>', container);
        assert.ok(container.innerHTML.includes('&lt;img'));
        assert.ok(!container.innerHTML.includes('<img onerror'));
    });

    it('parses CSV correctly', () => {
        const ctx = loadWithCSV();
        const ext = ctx.Extensions._registry.get('csv-table').def;

        // Simple
        assert.deepStrictEqual(ext._parseCSV('a,b\n1,2', ','), [['a','b'],['1','2']]);

        // Quoted with embedded comma
        assert.deepStrictEqual(ext._parseCSV('"hello, world",b\n1,2', ','), [['hello, world','b'],['1','2']]);

        // Escaped quotes
        assert.deepStrictEqual(ext._parseCSV('"a""b",c\n1,2', ','), [['a"b','c'],['1','2']]);

        // TSV
        assert.deepStrictEqual(ext._parseCSV('a\tb\n1\t2', '\t'), [['a','b'],['1','2']]);
    });
});

// ═══════════════════════════════════════════
// Diff Viewer Extension
// ═══════════════════════════════════════════

maybeDescribe('Diff Viewer extension', () => {
    function loadWithDiff() {
        const ctx = loadExtensions();
        ctx.Extensions.register({
            id: 'diff-viewer',
            _escapeHtml(str) {
                return str.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;');
            },
            init(extCtx) {
                const self = this;
                extCtx.renderers.register('diff-block', {
                    type: 'block',
                    priority: 10,
                    pattern: 'diff',
                    render(lang, code, container) {
                        const lines = code.split('\n');
                        let stats = { added: 0, removed: 0 };
                        const rendered = lines.map(line => {
                            const escaped = self._escapeHtml(line);
                            if (line.startsWith('+++ ') || line.startsWith('--- ')) {
                                return `<div class="diff-line diff-file-header">${escaped}</div>`;
                            }
                            if (line.startsWith('@@')) {
                                return `<div class="diff-line diff-hunk">${escaped}</div>`;
                            }
                            if (line.startsWith('+')) {
                                stats.added++;
                                return `<div class="diff-line diff-add">${escaped}</div>`;
                            }
                            if (line.startsWith('-')) {
                                stats.removed++;
                                return `<div class="diff-line diff-del">${escaped}</div>`;
                            }
                            return `<div class="diff-line diff-ctx">${escaped}</div>`;
                        }).join('');
                        container.innerHTML = `
                            <div class="diff-block ext-rendered">
                                <div class="diff-toolbar">
                                    <span class="diff-stat-add">+${stats.added}</span>
                                    <span class="diff-stat-del">-${stats.removed}</span>
                                </div>
                                <div class="diff-content">${rendered}</div>
                            </div>
                        `;
                    }
                });
            },
            destroy() {}
        });
        return ctx;
    }

    it('matches diff language', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        const handled = ctx.Extensions.runBlockRenderers('diff', '+added', container);
        assert.ok(handled);
    });

    it('does not match other languages', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        assert.strictEqual(ctx.Extensions.runBlockRenderers('python', 'x', { innerHTML: '' }), false);
    });

    it('classifies added lines', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('diff', '+new line', container);
        assert.ok(container.innerHTML.includes('diff-add'));
        assert.ok(container.innerHTML.includes('+1'));
    });

    it('classifies removed lines', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('diff', '-old line', container);
        assert.ok(container.innerHTML.includes('diff-del'));
        assert.ok(container.innerHTML.includes('-1'));
    });

    it('classifies hunk headers', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('diff', '@@ -1,3 +1,4 @@', container);
        assert.ok(container.innerHTML.includes('diff-hunk'));
    });

    it('classifies file headers', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('diff', '--- a/file.txt\n+++ b/file.txt', container);
        assert.ok(container.innerHTML.includes('diff-file-header'));
    });

    it('counts stats correctly', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('diff',
            '--- a/f.txt\n+++ b/f.txt\n@@ -1,2 +1,3 @@\n context\n+add1\n+add2\n-del1', container);
        assert.ok(container.innerHTML.includes('+2'));
        assert.ok(container.innerHTML.includes('-1'));
    });

    it('escapes HTML in diff content', async () => {
        const ctx = loadWithDiff();
        await ctx.Extensions.initAll();

        const container = { innerHTML: '' };
        ctx.Extensions.runBlockRenderers('diff', '+<script>alert(1)</script>', container);
        assert.ok(container.innerHTML.includes('&lt;script&gt;'));
    });
});

// ═══════════════════════════════════════════
// JS Sandbox Extension (Tool)
// ═══════════════════════════════════════════

maybeDescribe('JS Sandbox extension', () => {
    it('registers js_eval tool handler', async () => {
        const ctx = loadExtensions();
        ctx.Extensions.register({
            id: 'js-sandbox',
            init(extCtx) {
                extCtx.tools.handle('js_eval', async (args) => {
                    return { success: true, result: 'mock' };
                });
            },
            destroy() {}
        });
        await ctx.Extensions.initAll();
        assert.ok(ctx.Extensions._toolHandlers.has('js_eval'));
    });

    it('routes tool calls through the bridge', async () => {
        const ctx = loadExtensions();
        let receivedCode = null;

        ctx.Extensions.register({
            id: 'js-sandbox',
            init(extCtx) {
                extCtx.tools.handle('js_eval', async (args) => {
                    receivedCode = args.code;
                    return { success: true, result: '42', output: [] };
                });
            },
            destroy() {}
        });
        await ctx.Extensions.initAll();

        let resultEmitted = null;
        ctx.Events.on('tool.result.eval-1', (payload) => {
            resultEmitted = payload;
        });

        ctx.Events.emit('tool.call.eval-1', {
            call_id: 'eval-1',
            tool: 'js_eval',
            arguments: { code: '21 * 2' },
        }, { localOnly: true });

        await new Promise(r => setTimeout(r, 10));

        assert.strictEqual(receivedCode, '21 * 2');
        assert.ok(resultEmitted);
        const parsed = JSON.parse(resultEmitted.result);
        assert.strictEqual(parsed.success, true);
        assert.strictEqual(parsed.result, '42');
    });
});

// ═══════════════════════════════════════════
// Regex Tester Extension (Tool)
// ═══════════════════════════════════════════

maybeDescribe('Regex Tester extension', () => {
    function loadWithRegex() {
        const ctx = loadExtensions();
        ctx.Extensions.register({
            id: 'regex-tester',
            _formatMatch(match) {
                const result = {
                    fullMatch: match[0],
                    index: match.index,
                    length: match[0].length,
                };
                if (match.groups) {
                    result.namedGroups = {};
                    for (const [name, value] of Object.entries(match.groups)) {
                        result.namedGroups[name] = value;
                    }
                }
                if (match.length > 1) {
                    result.groups = [];
                    for (let i = 1; i < match.length; i++) {
                        result.groups.push(match[i] !== undefined ? match[i] : null);
                    }
                }
                return result;
            },
            _test(pattern, flags, inputs) {
                let regex;
                try {
                    regex = new RegExp(pattern, flags);
                } catch (e) {
                    return { success: false, error: `Invalid regex: ${e.message}`, pattern, flags };
                }
                const results = inputs.map(input => {
                    regex.lastIndex = 0;
                    const isGlobal = regex.global;
                    const matches = [];
                    if (isGlobal) {
                        let match;
                        while ((match = regex.exec(input)) !== null) {
                            matches.push(this._formatMatch(match));
                            if (match[0].length === 0) regex.lastIndex++;
                        }
                    } else {
                        const match = regex.exec(input);
                        if (match) matches.push(this._formatMatch(match));
                    }
                    return { input, matched: matches.length > 0, matchCount: matches.length, matches };
                });
                return {
                    success: true, pattern, flags,
                    summary: `${results.filter(r => r.matched).length}/${inputs.length} inputs matched`,
                    results,
                };
            },
            init(extCtx) {
                const self = this;
                extCtx.tools.handle('regex_test', async (args) => {
                    return self._test(args.pattern, args.flags || '', args.inputs || []);
                });
            },
            destroy() {}
        });
        return ctx;
    }

    it('registers regex_test tool handler', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();
        assert.ok(ctx.Extensions._toolHandlers.has('regex_test'));
    });

    it('matches simple patterns', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();
        const ext = ctx.Extensions._registry.get('regex-tester').def;

        const result = ext._test('^\\d+$', '', ['123', 'abc', '456']);
        assert.strictEqual(result.success, true);
        assert.strictEqual(result.summary, '2/3 inputs matched');
        assert.strictEqual(result.results[0].matched, true);
        assert.strictEqual(result.results[1].matched, false);
        assert.strictEqual(result.results[2].matched, true);
    });

    it('captures groups', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();
        const ext = ctx.Extensions._registry.get('regex-tester').def;

        const result = ext._test('(\\d{3})-(\\d{4})', '', ['555-1234']);
        assert.strictEqual(result.success, true);
        assert.strictEqual(result.results[0].matches[0].fullMatch, '555-1234');
        assert.deepStrictEqual(result.results[0].matches[0].groups, ['555', '1234']);
    });

    it('handles global flag with multiple matches', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();
        const ext = ctx.Extensions._registry.get('regex-tester').def;

        const result = ext._test('\\d+', 'g', ['a1b2c3']);
        assert.strictEqual(result.results[0].matchCount, 3);
        assert.strictEqual(result.results[0].matches[0].fullMatch, '1');
        assert.strictEqual(result.results[0].matches[1].fullMatch, '2');
        assert.strictEqual(result.results[0].matches[2].fullMatch, '3');
    });

    it('returns error for invalid patterns', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();
        const ext = ctx.Extensions._registry.get('regex-tester').def;

        const result = ext._test('[invalid', '', ['test']);
        assert.strictEqual(result.success, false);
        assert.ok(result.error.includes('Invalid regex'));
    });

    it('handles named groups', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();
        const ext = ctx.Extensions._registry.get('regex-tester').def;

        const result = ext._test('(?<year>\\d{4})-(?<month>\\d{2})', '', ['2025-02']);
        assert.strictEqual(result.results[0].matches[0].namedGroups.year, '2025');
        assert.strictEqual(result.results[0].matches[0].namedGroups.month, '02');
    });

    it('routes through tool bridge', async () => {
        const ctx = loadWithRegex();
        await ctx.Extensions.initAll();

        let resultEmitted = null;
        ctx.Events.on('tool.result.regex-1', (payload) => {
            resultEmitted = payload;
        });

        ctx.Events.emit('tool.call.regex-1', {
            call_id: 'regex-1',
            tool: 'regex_test',
            arguments: { pattern: '^hello', flags: 'i', inputs: ['Hello world', 'goodbye'] },
        }, { localOnly: true });

        await new Promise(r => setTimeout(r, 10));

        assert.ok(resultEmitted);
        const parsed = JSON.parse(resultEmitted.result);
        assert.strictEqual(parsed.success, true);
        assert.strictEqual(parsed.summary, '1/2 inputs matched');
    });
});