import React, { createContext, useState, use, useEffect, Suspense } from "react";
import { Theme } from "@code-hike/lighter";
import type { HighlightedCode as CodeHikeHighlightedCode } from "codehike/code";

// Client highlight dispatch: routes to xyd's own Rust (WASM) highlighter when the
// `engine.highlighter: "rust"` toggle is on and a highlighter is registered, else
// codehike (default). Same signature as codehike's highlight().
import { highlight } from "../highlightDispatch";

import defaultTheme from "../themes/cosmo"
import { Loader } from "../../kit";

export interface CodeThemeProps {
    codeblocks?: CodeThemeBlockProps[];
    theme?: Theme
    children: React.ReactNode;
}

export interface SyntaxHighlightedCode extends CodeHikeHighlightedCode {
    title?: string
}

export interface CodeThemeBlockProps {
    /** This is the raw code. May include annotation comments. */
    value: string;
    /** The programming language. */
    lang: string;
    /** Metadata string (the content after the language name in a markdown codeblock). */
    meta: string;

    title?: string

    /** The highlighted code. */
    highlighted?: SyntaxHighlightedCode
}

const CodeThemeProvider = createContext<{
    highlighted: SyntaxHighlightedCode[];
}>({
    highlighted: [],
});

export function useCodeTheme() {
    return use(CodeThemeProvider);
}

// A module-level cache of highlighted code, keyed by theme + language + source.
// Virtualized lists unmount off-screen code samples; without this, scrolling a
// sample back into view re-runs codehike `highlight()` (async → a visible
// flash). With it, a re-mounted sample reuses its tokens synchronously — the
// state is seeded from the cache on mount, so a cache hit renders instantly.
const HIGHLIGHT_CACHE_MAX = 3000;
const highlightCache = new Map<string, SyntaxHighlightedCode>();

function cacheKey(codeblock: CodeThemeBlockProps, theme: Theme): string {
    const t = theme as any;
    const themeKey = typeof t === "string" ? t : (t?.name ?? "t");
    return `${themeKey} ${codeblock.lang} ${codeblock.value}`;
}

function cacheGet(codeblock: CodeThemeBlockProps, theme: Theme) {
    return highlightCache.get(cacheKey(codeblock, theme));
}

function cacheSet(codeblock: CodeThemeBlockProps, theme: Theme, value: SyntaxHighlightedCode) {
    if (highlightCache.size >= HIGHLIGHT_CACHE_MAX) {
        // Evict the oldest entry — Map preserves insertion order.
        const oldest = highlightCache.keys().next().value;
        if (oldest !== undefined) highlightCache.delete(oldest);
    }
    highlightCache.set(cacheKey(codeblock, theme), value);
}

function hasHighlight(codeblock: CodeThemeBlockProps, theme: Theme): boolean {
    return Boolean(
        (codeblock.highlighted && codeblock.highlighted.tokens) ||
        cacheGet(codeblock, theme),
    );
}

// The highlighted codeblocks we already have (server-provided or cached) — used
// to seed state synchronously so a cache hit renders with no loader flash.
function initializeHighlighted(codeblocks: CodeThemeBlockProps[] | undefined, theme: Theme): SyntaxHighlightedCode[] {
    if (!codeblocks) return [];

    return codeblocks.map((codeblock) => {
        if (codeblock.highlighted && codeblock.highlighted.tokens) {
            return {
                ...codeblock.highlighted,
                meta: codeblock.highlighted?.meta || codeblock.meta,
                title: codeblock.title,
            };
        }

        const cached = cacheGet(codeblock, theme);
        if (cached) {
            return { ...cached, meta: cached.meta || codeblock.meta, title: codeblock.title };
        }

        return null;
    }).filter(Boolean) as SyntaxHighlightedCode[];
}

// TODO: lazy-loading - some codeblocks on server and another on client
// TODO: lazdy client loading only on specific codeblock
export function CodeTheme(props: CodeThemeProps) {
    const theme = props.theme || defaultTheme;

    // Seed from server-highlighted + already-cached codeblocks.
    const [highlighted, setHighlighted] = useState<SyntaxHighlightedCode[] | undefined>(
        () => initializeHighlighted(props.codeblocks, theme),
    );

    useEffect(() => {
        setHighlighted(initializeHighlighted(props.codeblocks, theme));
    }, [props.codeblocks, theme]);

    useEffect(() => {
        if (!props.codeblocks) {
            return;
        }

        // Only fetch for codeblocks we don't already have. On a full cache hit
        // (a virtualized sample scrolled back into view) this is a no-op.
        const needsHighlighting = props.codeblocks.some(
            (codeblock) => !hasHighlight(codeblock, theme),
        );

        if (needsHighlighting) {
            clientSideHighlight();
        }
    }, [props.codeblocks, theme]);

    async function clientSideHighlight() {
        if (!props.codeblocks) {
            return;
        }

        const newHighlighted = await fetchHighlight(props.codeblocks, theme);
        setHighlighted(newHighlighted);
    }

    const withTheme = <CodeThemeProvider
        value={{
            highlighted: highlighted || [],
        }}
    >
        {props.children}
    </CodeThemeProvider>

    // Ready once every codeblock has highlighted tokens (server-provided, cached,
    // or freshly fetched). Only then render the code — otherwise a brief loader.
    const total = props.codeblocks?.length ?? 0;
    const ready = total === 0 || (highlighted?.length ?? 0) >= total;

    return ready ? withTheme : <Loader />;
}

// TODO: own server with grammars (codehike)
async function fetchHighlight(codeblocks: CodeThemeBlockProps[], theme: Theme) {
    return await Promise.all(
        codeblocks?.map(async (codeblock) => {
            const cached = cacheGet(codeblock, theme);
            if (cached) return cached;

            const result = await highlight(codeblock, theme);
            cacheSet(codeblock, theme, result);
            return result;
        }),
    );
}

/**
 * Pre-highlight codeblocks INTO the shared cache without rendering them, so a
 * consumer (e.g. a virtualized list) can warm the samples just outside the
 * viewport — when they scroll in, {@link CodeTheme} finds them cached and paints
 * instantly. Idempotent: already-cached codeblocks are skipped. Never throws.
 */
export async function prewarmHighlight(
    codeblocks: CodeThemeBlockProps[] | undefined,
    theme?: Theme,
): Promise<void> {
    if (!codeblocks?.length) return;

    const t = theme || defaultTheme;
    await Promise.all(
        codeblocks.map(async (codeblock) => {
            if (!codeblock?.value || cacheGet(codeblock, t)) return;
            try {
                const result = await highlight(codeblock, t);
                cacheSet(codeblock, t, result);
            } catch {
                // A failed prewarm just means it highlights on mount instead.
            }
        }),
    );
}