import { useState, useEffect, useMemo } from "react";
import type { ThemedToken } from "shiki";

/* ── Lazy singleton highlighter ─────────────────────────── */

let highlighterPromise: Promise<Awaited<
  ReturnType<typeof import("shiki")["createHighlighter"]>
> | null> | null = null;

const LANGS = [
  "typescript",
  "javascript",
  "tsx",
  "jsx",
  "json",
  "css",
  "html",
  "yaml",
  "markdown",
  "bash",
  "python",
  "go",
  "rust",
  "sql",
  "toml",
  "xml",
] as const;

function getHighlighter() {
  if (!highlighterPromise) {
    highlighterPromise = import("shiki")
      .then(({ createHighlighter }) =>
        createHighlighter({
          themes: ["github-dark-default"],
          langs: [...LANGS],
        }),
      )
      .catch(() => null);
  }
  return highlighterPromise;
}

// Pre-warm on module load (non-blocking)
getHighlighter();

/* ── Language detection ──────────────────────────────────── */

const EXT_TO_LANG: Record<string, string> = {
  ts: "typescript",
  tsx: "tsx",
  js: "javascript",
  jsx: "jsx",
  mjs: "javascript",
  cjs: "javascript",
  json: "json",
  css: "css",
  scss: "css",
  html: "html",
  htm: "html",
  yml: "yaml",
  yaml: "yaml",
  md: "markdown",
  mdx: "markdown",
  sh: "bash",
  bash: "bash",
  zsh: "bash",
  py: "python",
  go: "go",
  rs: "rust",
  sql: "sql",
  toml: "toml",
  xml: "xml",
};

function detectLang(path: string): string | null {
  const ext = path.split(".").pop()?.toLowerCase() ?? "";
  return EXT_TO_LANG[ext] ?? null;
}

/* ── Types ───────────────────────────────────────────────── */

export type TokenizedLine = ThemedToken[];
/** Maps newLineNumber → highlighted tokens for that line */
export type LineTokenMap = Map<number, TokenizedLine>;

interface DiffLineInput {
  content: string;
  newLineNumber: number | null;
  type: "added" | "removed" | "context";
}

/* ── Hook ────────────────────────────────────────────────── */

/**
 * Highlights the "new-side" content of a file diff.
 * Returns null until highlighting is ready (progressive enhancement).
 * Only context + added lines are highlighted (removed lines fall back to plain text).
 */
export function useHighlightedFile(
  filePath: string,
  allLines: DiffLineInput[],
): LineTokenMap | null {
  const [tokenMap, setTokenMap] = useState<LineTokenMap | null>(null);

  // Build stable code string + line number mapping
  const { code, lineNums, cacheKey } = useMemo(() => {
    const entries: { lineNum: number; content: string }[] = [];
    for (const line of allLines) {
      if (
        line.newLineNumber !== null &&
        (line.type === "context" || line.type === "added")
      ) {
        entries.push({ lineNum: line.newLineNumber, content: line.content });
      }
    }
    entries.sort((a, b) => a.lineNum - b.lineNum);
    const c = entries.map((e) => e.content).join("\n");
    return {
      code: c,
      lineNums: entries.map((e) => e.lineNum),
      cacheKey: `${filePath}:${c.length}:${entries.length}`,
    };
  }, [filePath, allLines]);

  useEffect(() => {
    const lang = detectLang(filePath);
    if (!lang || !code) {
      setTokenMap(null);
      return;
    }

    let cancelled = false;

    getHighlighter().then((highlighter) => {
      if (cancelled || !highlighter) return;

      try {
        const result = highlighter.codeToTokens(code, {
          lang: lang as Parameters<typeof highlighter.codeToTokens>[1]["lang"],
          theme: "github-dark-default",
        });
        const map: LineTokenMap = new Map();

        result.tokens.forEach((lineTokens: ThemedToken[], idx: number) => {
          if (idx < lineNums.length) {
            map.set(lineNums[idx], lineTokens);
          }
        });

        if (!cancelled) setTokenMap(map);
      } catch {
        // Language not loaded or parse error — no highlighting
      }
    });

    return () => {
      cancelled = true;
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [cacheKey]);

  return tokenMap;
}
