import type { MarkdownDocument, Node } from 'comark';
import { createSerializedMarkdownParser as createComarkParser } from 'comark';
import breaks from 'comark/plugins/breaks';
import footnotes from 'comark/plugins/footnotes';
import math from 'comark/plugins/math';
import { citationPlugin } from './citations';

export type MarkdownTree = MarkdownDocument;
export type MarkdownAstNode = Node;
export type MarkdownElementNode = Exclude<MarkdownAstNode, string | [null, ...unknown[]]>;

const footnoteSyntax = {
    name: 'footnote-syntax',
    markdownItPlugins: [
        (markdown: any) => markdown.inline.ruler.before('link', 'footnote_inline', (state: any, silent: boolean) => {
            const start = state.pos;
            if (!state.src.startsWith('[^', start)) {
                return false;
            }

            const end = state.src.indexOf(']', start + 2);
            const label = end === -1 ? '' : state.src.slice(start + 2, end);
            if (end === -1 || label.length === 0 || /\s/.test(label) || silent) {
                return false;
            }

            state.push('mdc_inline_span', 'span', 1);
            state.push('text', '', 0).content = `^${label}`;
            state.push('mdc_inline_span', 'span', -1);
            state.pos = end + 1;
            return true;
        }),
    ],
};

const parserPlugins = [
    breaks(),
    footnoteSyntax,
    footnotes(),
    math(),
] as const;

const createConfiguredParser = (autoClose: boolean, citations: boolean) => createComarkParser({
    registerDefaultPlugins: false,
    autoClose,
    plugins: citations ? [...parserPlugins, citationPlugin] : parserPlugins,
});

const protectDollarRunsInLine = (line: string): string => {
    let result = '';
    let segment = '';
    let codeDelimiterLength = 0;

    const flushSegment = () => {
        result += segment.replace(/\${3,}/g, (run) => run.replaceAll('$', '&#36;'));
        segment = '';
    };

    for (let index = 0; index < line.length;) {
        let delimiterLength = 0;
        if (line[index] === '`') {
            while (line[index + delimiterLength] === '`') {
                delimiterLength += 1;
            }
        }

        if (codeDelimiterLength !== 0) {
            result += line.slice(index, index + (delimiterLength || 1));
            if (delimiterLength === codeDelimiterLength) {
                codeDelimiterLength = 0;
            }
            index += delimiterLength || 1;
            continue;
        }

        if (delimiterLength === 0) {
            segment += line[index];
            index += 1;
            continue;
        }

        flushSegment();
        result += line.slice(index, index + delimiterLength);
        codeDelimiterLength = delimiterLength;
        index += delimiterLength;
    }

    flushSegment();
    return result;
};

const protectDollarRuns = (content: string): string => {
    // Fast path: most messages contain no dollar runs; avoid line-splitting entirely.
    if (!content.includes('$')) {
        return content;
    }
    const lines = content.split('\n');
    let inFence = false;

    return lines.map((line) => {
        const fence = /^\s{0,3}(`{3,}|~{3,})/.exec(line);
        if (fence !== null) {
            inFence = !inFence;
            return line;
        }

        if (inFence) {
            return line;
        }

        return protectDollarRunsInLine(line);
    }).join('\n');
};

export const createMarkdownParser = (options: { citations?: boolean } = {}) => {
    const citations = options.citations === true;
    const streamingParser = createConfiguredParser(true, citations);
    const finalParser = createConfiguredParser(false, citations);

    return (content: string, options: { streaming?: boolean } = {}) => {
        const protectedContent = protectDollarRuns(content);
        // Do not let auto-close complete an unfinished citation on the model's behalf.
        const incompleteCitation = citations && /\[\[cite:[a-zA-Z0-9_,\-]*\]?$/.test(protectedContent);
        if (options.streaming === true && !incompleteCitation) {
            return streamingParser(protectedContent, { streaming: true });
        }

        return finalParser(protectedContent, { streaming: false });
    };
};

export const agentParser = createMarkdownParser({ citations: true });
export const userParser = createMarkdownParser({ citations: false });

export const getNodeTag = (node: MarkdownAstNode): string | null => {
    if (!Array.isArray(node) || node[0] === null) {
        return null;
    }

    return node[0];
};

export const getNodeAttributes = (node: MarkdownAstNode): Record<string, unknown> => {
    if (!Array.isArray(node) || node[0] === null || typeof node[1] !== 'object' || node[1] === null) {
        return {};
    }

    return node[1] as Record<string, unknown>;
};

export const getNodeChildren = (node: MarkdownAstNode): MarkdownAstNode[] => {
    if (!Array.isArray(node) || node[0] === null) {
        return [];
    }

    return node.slice(2) as MarkdownAstNode[];
};

export const textContent = (node: MarkdownAstNode): string => {
    if (typeof node === 'string') {
        return node;
    }

    if (!Array.isArray(node) || node[0] === null) {
        return '';
    }

    return getNodeChildren(node).map(textContent).join('');
};

const splitParagraph = (attributes: Record<string, unknown>, children: MarkdownAstNode[]): MarkdownElementNode[] => {
    const paragraphs: MarkdownElementNode[] = [];
    let current: MarkdownAstNode[] = [];

    for (const child of children) {
        if (getNodeTag(child) === 'br') {
            paragraphs.push(['p', { ...attributes }, ...current] as MarkdownElementNode);
            current = [];
            continue;
        }

        current.push(child);
    }

    if (current.length > 0 || paragraphs.length === 0) {
        paragraphs.push(['p', { ...attributes }, ...current] as MarkdownElementNode);
    }

    return paragraphs;
};

const hasBrChild = (children: MarkdownAstNode[]): boolean =>
    children.some((child) => getNodeTag(child) === 'br');

/**
 * Transform a single element node into its replacement list. Returns a
 * single-element array holding the ORIGINAL node reference when nothing
 * changed, so callers can skip all allocation for unaffected subtrees.
 */
const transformElement = (node: MarkdownElementNode): MarkdownAstNode[] => {
    const tag = node[0];
    const children = getNodeChildren(node);
    const transformedChildren = transformChildren(children);

    if (tag === 'p') {
        if (!hasBrChild(transformedChildren)) {
            if (transformedChildren === children) return [node];
            return [[tag, { ...node[1] }, ...transformedChildren] as MarkdownElementNode];
        }
        return splitParagraph({ ...node[1] }, transformedChildren);
    }

    if ((tag === 'li' || tag === 'blockquote') && hasBrChild(transformedChildren)) {
        const segments: MarkdownAstNode[][] = [[]];
        for (const child of transformedChildren) {
            if (getNodeTag(child) === 'br') {
                segments.push([]);
            } else {
                segments[segments.length - 1]?.push(child);
            }
        }

        return [[
            tag,
            { ...node[1] },
            ...segments
                .filter((segment) => segment.length > 0)
                .map((segment) => ['p', {}, ...segment] as MarkdownElementNode),
        ] as MarkdownElementNode];
    }

    if (transformedChildren === children) return [node];
    return [[tag, { ...node[1] }, ...transformedChildren] as MarkdownElementNode];
};

const transformChildren = (nodes: MarkdownAstNode[]): MarkdownAstNode[] => {
    let result: MarkdownAstNode[] | null = null;

    for (let index = 0; index < nodes.length; index++) {
        const node = nodes[index]!;
        const replacement = typeof node === 'string' || !Array.isArray(node) || node[0] === null
            ? [node]
            : transformElement(node);

        if (result === null) {
            if (replacement.length === 1 && replacement[0] === node) continue;
            result = nodes.slice(0, index);
        }
        result.push(...replacement);
    }

    return result ?? nodes;
};

/**
 * Comark's breaks plugin emits `br` nodes. Promote each hard break to a
 * sibling paragraph to keep the chat renderer's established layout.
 *
 * Copy-on-write: trees without `br` nodes are returned by reference with no
 * allocation, so the common case costs a single tag scan per node.
 */
export const splitParagraphBreaks = (tree: MarkdownTree): MarkdownTree => {
    const nodes = transformChildren(tree.nodes);
    return nodes === tree.nodes ? tree : { ...tree, nodes };
};

export const toVueAttributes = (attributes: Record<string, unknown>): Record<string, unknown> => {
    const result: Record<string, unknown> = {};

    for (const [key, value] of Object.entries(attributes)) {
        if (key !== '$') {
            result[key] = value;
        }
    }

    return result;
};