import { useState, useEffect, useMemo } from "react";
import type { ThemedToken } from "shiki";
import type { HighlightRequest, HighlightResponse } from "./highlight-worker";

/* ── Worker pool (module-level, shared across all hook instances) ─────── */

type PendingResolve = (response: HighlightResponse) => void;

let workers: Worker[] = [];
let requestCount = 0;
const MAX_WORKERS = 2;
const pending = new Map<string, PendingResolve>();

let workersInitialized = false;

function initWorkers(): void {
  if (workersInitialized) return;
  workersInitialized = true;
  try {
    workers = Array.from({ length: MAX_WORKERS }, () => {
      const w = new Worker(
        new URL("./highlight-worker.ts", import.meta.url),
        { type: "module" },
      );
      w.addEventListener("message", (event: MessageEvent<HighlightResponse>) => {
        const resolve = pending.get(event.data.id);
        if (resolve) {
          pending.delete(event.data.id);
          resolve(event.data);
        }
      });
      return w;
    });
  } catch {
    // CSP or browser compat — fall back to chunked main-thread highlighting
    workers = [];
  }
}

function highlightWithWorker(
  id: string,
  language: string,
  code: string,
  lineNumbers: number[],
  filePath: string,
): Promise<HighlightResponse> {
  return new Promise<HighlightResponse>((resolve) => {
    pending.set(id, resolve);
    const worker = workers[requestCount % MAX_WORKERS];
    requestCount++;
    const req: HighlightRequest = { id, filePath, language, code, lineNumbers };
    worker.postMessage(req);
  });
}

/* ── Lazy singleton highlighter (for main-thread fallback) ───────────── */

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;
}

/* ── Chunked main-thread fallback ────────────────────────────────────── */

async function highlightChunked(
  code: string,
  language: string,
  lineNumbers: number[],
  signal: AbortSignal,
): Promise<LineTokenMap> {
  const CHUNK = 200;
  const result: LineTokenMap = new Map();
  const lines = code.split("\n");
  const highlighter = await getHighlighter();
  if (!highlighter) return result;

  for (let i = 0; i < lines.length; i += CHUNK) {
    if (signal.aborted) break;
    const chunkLines = lines.slice(i, i + CHUNK);
    const chunkCode = chunkLines.join("\n");
    try {
      const tokenized = highlighter.codeToTokens(chunkCode, {
        lang: language as Parameters<typeof highlighter.codeToTokens>[1]["lang"],
        theme: "github-dark-default",
      });
      tokenized.tokens.forEach((lineTokens: ThemedToken[], idx: number) => {
        const lineNum = lineNumbers[i + idx];
        if (lineNum !== undefined) result.set(lineNum, lineTokens);
      });
    } catch {
      // Skip unparseable chunk
    }

    // Yield between chunks to avoid blocking the main thread
    await new Promise<void>((r) => {
      if (
        "scheduler" in globalThis &&
        "yield" in (globalThis as Record<string, unknown>).scheduler
      ) {
        (
          (globalThis as Record<string, unknown>).scheduler as {
            yield: () => Promise<void>;
          }
        )
          .yield()
          .then(r);
      } else {
        setTimeout(r, 0);
      }
    });
  }
  return result;
}

/* ── 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, returning a map of
 * line number → syntax tokens.
 *
 * Progressive rendering: returns `null` while highlighting is in progress.
 * Callers (HunkRows → LineWithComments) render plain text when `null` and
 * patch in highlighted tokens on re-render once the worker or chunked call
 * resolves.
 *
 * Worker path: uses a module-level pool of 2 Web Workers. Round-robin
 * assignment. Late responses after unmount are silently discarded.
 *
 * Fallback path: if Worker construction fails (CSP, browser compat),
 * falls back to chunked main-thread highlighting via codeToTokens (200
 * lines/chunk) with scheduler.yield()/setTimeout(0) between chunks.
 */
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;
    }

    initWorkers(); // no-op after first call

    const id = crypto.randomUUID();
    let unmounted = false;
    const abortController = new AbortController();

    if (workers.length > 0) {
      highlightWithWorker(id, lang, code, lineNums, filePath).then((response) => {
        if (unmounted) return; // ignore late responses after unmount
        if (response.error || response.tokens.length === 0) {
          setTokenMap(null);
          return;
        }
        const map: LineTokenMap = new Map();
        for (const { lineNumber, tokens } of response.tokens) {
          map.set(lineNumber, tokens);
        }
        setTokenMap(map);
      });
    } else {
      highlightChunked(code, lang, lineNums, abortController.signal).then((map) => {
        if (unmounted) return;
        setTokenMap(map.size > 0 ? map : null);
      });
    }

    return () => {
      unmounted = true;
      abortController.abort();
      // Remove pending resolver so a late worker response is silently dropped
      pending.delete(id);
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [cacheKey]);

  return tokenMap;
}
