"use client";
import * as React from "react";
import ReactMarkdown, { type Components } from "react-markdown";
import remarkGfm from "remark-gfm";
import remarkMath from "remark-math";
import rehypeKatex from "rehype-katex";
import "katex/dist/katex.min.css";
import { CodeBlock } from "./code-block";
import { normalizeMath, splitStreamingMarkdown } from "@/lib/chat/markdown-blocks";
import { cn } from "@/lib/utils";

/**
 * Chat Markdown renderer.
 *
 * - GFM (tables, task lists, strikethrough) + LaTeX (`$…$`, `$$…$$`, and the `\(…\)` / `\[…\]` forms
 *   models emit, normalised in `lib/chat/markdown-blocks`). Raw HTML is skipped (never injected).
 * - Streaming: the text is split into completed top-level blocks that are rendered once and memoised;
 *   only the trailing block re-parses on each token, so long answers stay cheap and tables/code do not flicker.
 */
const components = (wrap: boolean): Components => ({
  code({ className, children, ...props }) {
    const match = /language-(\w[\w+#-]*)/.exec(className ?? "");
    const text = String(children ?? "").replace(/\n$/, "");
    const isBlock = Boolean(match) || text.includes("\n");
    if (isBlock) return <CodeBlock code={text} lang={match?.[1]} wrap={wrap} />;
    return (
      <code className={className} {...props}>
        {children}
      </code>
    );
  },
  pre({ children }) {
    return <>{children}</>;
  },
  a({ href, children }) {
    return (
      <a href={href} target="_blank" rel="noopener noreferrer nofollow">
        {children}
      </a>
    );
  },
  table({ children }) {
    return (
      <div className="my-2 overflow-x-auto scrollbar-thin">
        <table>{children}</table>
      </div>
    );
  },
});

const REMARK_PLUGINS = [remarkGfm, remarkMath];
const REHYPE_PLUGINS = [[rehypeKatex, { output: "html", throwOnError: false, strict: false }]] as never[];

const Block = React.memo(function Block({ content, comps }: { content: string; comps: Components }) {
  return (
    <ReactMarkdown remarkPlugins={REMARK_PLUGINS} rehypePlugins={REHYPE_PLUGINS} components={comps} skipHtml>
      {content}
    </ReactMarkdown>
  );
});

export const Markdown = React.memo(function Markdown({ content, className, wrap = false, streaming = false }: { content: string; className?: string; wrap?: boolean; streaming?: boolean }) {
  const comps = React.useMemo(() => components(wrap), [wrap]);
  const normalized = React.useMemo(() => normalizeMath(content), [content]);
  const split = React.useMemo(() => (streaming ? splitStreamingMarkdown(normalized, { minTailChars: 24 }) : null), [normalized, streaming]);
  return (
    <div className={cn("prose-chat", wrap && "wrap-code", streaming && "caret", className)}>
      {split ? (
        <>
          {split.blocks.map((b, i) => (
            <Block key={i} content={b} comps={comps} />
          ))}
          {split.tail ? <Block content={split.tail} comps={comps} /> : null}
        </>
      ) : (
        <Block content={normalized} comps={comps} />
      )}
    </div>
  );
});