"use client";
import * as React from "react";
import { Check, ChevronsDownUp, ChevronsUpDown, Copy, Download, WrapText } from "lucide-react";
import { extensionForLanguage } from "@/lib/chat/markdown-blocks";
import { useDebounced } from "@/lib/client/hooks";
import { cn } from "@/lib/utils";

type Highlighter = { codeToHtml: (code: string, opts: { lang: string; themes: { light: string; dark: string }; defaultColor: false | "light" }) => string; getLoadedLanguages: () => string[]; loadLanguage: (lang: never) => Promise<void> };

let highlighterPromise: Promise<Highlighter> | null = null;
const KNOWN = new Set<string>();

async function getHighlighter(): Promise<Highlighter> {
  if (!highlighterPromise) {
    highlighterPromise = (async () => {
      const { createHighlighter } = await import("shiki/bundle/web");
      const { createJavaScriptRegexEngine } = await import("shiki/engine/javascript");
      const hl = await createHighlighter({ themes: ["github-light-default", "github-dark-default"], langs: ["javascript", "typescript", "tsx", "jsx", "json", "bash", "shell", "python", "markdown", "html", "css", "sql", "yaml"], engine: createJavaScriptRegexEngine({ forgiving: true }) });
      for (const l of hl.getLoadedLanguages()) KNOWN.add(l);
      return hl as unknown as Highlighter;
    })();
  }
  return highlighterPromise;
}

const ALIASES: Record<string, string> = { js: "javascript", ts: "typescript", sh: "bash", zsh: "bash", shell: "bash", py: "python", yml: "yaml", md: "markdown", rb: "ruby", rs: "rust", golang: "go", "c++": "cpp", cs: "csharp", kt: "kotlin", plaintext: "text", txt: "text" };

/** Lines above which a block starts collapsed (Expand reveals everything). */
const COLLAPSE_LINES = 28;

export function CodeBlock({ code, lang, wrap }: { code: string; lang?: string; wrap?: boolean }) {
  const [html, setHtml] = React.useState<string | null>(null);
  const [copied, setCopied] = React.useState(false);
  const [localWrap, setLocalWrap] = React.useState<boolean | null>(null);
  const [expanded, setExpanded] = React.useState(false);
  const language = React.useMemo(() => {
    const l = (lang ?? "").toLowerCase().trim();
    return ALIASES[l] ?? l;
  }, [lang]);
  const effectiveWrap = localWrap ?? wrap ?? false;
  // While streaming the block changes on every token; highlight the settled version only.
  const settled = useDebounced(code, 120);
  const lineCount = React.useMemo(() => (code ? code.split("\n").length : 0), [code]);
  const collapsible = lineCount > COLLAPSE_LINES;

  React.useEffect(() => {
    let cancelled = false;
    if (!language || language === "text") {
      // eslint-disable-next-line react-hooks/set-state-in-effect
      setHtml(null);
      return;
    }
    (async () => {
      try {
        const hl = await getHighlighter();
        if (!KNOWN.has(language)) {
          try {
            await hl.loadLanguage(language as never);
            KNOWN.add(language);
          } catch {
            if (!cancelled) setHtml(null);
            return;
          }
        }
        const out = hl.codeToHtml(settled, { lang: language, themes: { light: "github-light-default", dark: "github-dark-default" }, defaultColor: "light" });
        if (!cancelled) setHtml(out);
      } catch {
        if (!cancelled) setHtml(null);
      }
    })();
    return () => {
      cancelled = true;
    };
  }, [settled, language]);

  const copy = async () => {
    try {
      await navigator.clipboard.writeText(code);
      setCopied(true);
      setTimeout(() => setCopied(false), 1400);
    } catch {
      /* ignore */
    }
  };

  const download = () => {
    const ext = extensionForLanguage(language);
    const blob = new Blob([code], { type: "text/plain;charset=utf-8" });
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = ext === "Dockerfile" ? "Dockerfile" : `snippet.${ext}`;
    document.body.appendChild(a);
    a.click();
    a.remove();
    setTimeout(() => URL.revokeObjectURL(url), 1000);
  };

  // Show the highlighted HTML only when it matches the current code (avoids stale colouring while streaming).
  const showHtml = html && settled === code;

  return (
    <div className="group/code my-2 overflow-hidden rounded-lg border border-border bg-bg-subtle text-[13px]">
      <div className="flex h-9 items-center justify-between gap-2 border-b border-border px-2.5 sm:h-8 sm:px-3">
        <span className="truncate font-mono text-[11px] text-fg-subtle">
          {language || "text"}
          {lineCount > 1 ? <span className="hidden sm:inline"> · {lineCount} lines</span> : null}
        </span>
        <div className="flex items-center gap-0.5">
          {collapsible ? (
            <ToolBtn label={expanded ? "Collapse" : "Expand"} onClick={() => setExpanded((e) => !e)}>
              {expanded ? <ChevronsDownUp className="size-3.5" /> : <ChevronsUpDown className="size-3.5" />}
            </ToolBtn>
          ) : null}
          <ToolBtn label={effectiveWrap ? "Disable line wrap" : "Wrap long lines"} onClick={() => setLocalWrap((w) => !(w ?? wrap ?? false))} active={effectiveWrap}>
            <WrapText className="size-3.5" />
          </ToolBtn>
          <ToolBtn label="Download" onClick={download}>
            <Download className="size-3.5" />
          </ToolBtn>
          <button onClick={copy} className="tap inline-flex h-7 items-center gap-1 rounded px-1.5 text-[11px] text-fg-subtle hover:bg-bg-muted hover:text-fg" aria-label="Copy code">
            {copied ? <Check className="size-3.5 text-success" /> : <Copy className="size-3.5" />}
            <span className="hidden sm:inline">{copied ? "Copied" : "Copy"}</span>
          </button>
        </div>
      </div>
      <div className={cn("relative", collapsible && !expanded && "max-h-[420px] overflow-hidden")}>
        <div className={cn("overflow-x-auto scrollbar-thin", effectiveWrap && "[&_pre]:whitespace-pre-wrap [&_pre]:break-words [&_code]:whitespace-pre-wrap")}>
          {showHtml ? (
            <div className="[&_pre]:m-0 [&_pre]:bg-transparent! [&_pre]:p-3 [&_pre]:leading-[1.55] [&_code]:font-mono [&_code]:text-[12.5px]" dangerouslySetInnerHTML={{ __html: html }} />
          ) : (
            <pre className="m-0 p-3 leading-[1.55]">
              <code className="font-mono text-[12.5px]">{code}</code>
            </pre>
          )}
        </div>
        {collapsible && !expanded ? (
          <button type="button" onClick={() => setExpanded(true)} className="absolute inset-x-0 bottom-0 flex h-16 items-end justify-center bg-gradient-to-t from-bg-subtle via-bg-subtle/80 to-transparent pb-2 text-[12px] font-medium text-fg-muted hover:text-fg">
            Show all {lineCount} lines
          </button>
        ) : null}
      </div>
    </div>
  );
}

function ToolBtn({ label, onClick, active, children }: { label: string; onClick: () => void; active?: boolean; children: React.ReactNode }) {
  return (
    <button type="button" onClick={onClick} className={cn("tap rounded p-1.5 text-fg-subtle hover:bg-bg-muted hover:text-fg", active && "text-accent")} aria-label={label} title={label}>
      {children}
    </button>
  );
}