import { createElement, useMemo, type ReactNode } from "react";
import { cn } from "../../lib/cn";
import type { AIInputMention } from "../inputs/ai-input";
import {
  ChecklistBlock,
  CodeBlock,
  DiagramBlock,
  EmbedBlock,
  FileBlock,
  KeyValueBlock,
  MathBlock,
  MediaBlock,
  QuoteBlock,
  StepsBlock,
  TableBlock,
  TimelineBlock,
} from "./advanced-blocks";

export type MarkdownRenderContext = {
  language?: string;
  value: string;
};

export type MarkdownProps = {
  allowRawHtml?: boolean;
  className?: string;
  onCopy?: (value: string) => void;
  mentions?: AIInputMention[];
  mentionHighlightClassName?: string;
  renderCode?: (code: string, context: MarkdownRenderContext) => ReactNode;
  renderDiagram?: (source: string, context: MarkdownRenderContext) => ReactNode;
  renderImage?: (src: string, alt: string) => ReactNode;
  renderLink?: (href: string, label: string) => ReactNode;
  renderMath?: (source: string, display: boolean) => ReactNode;
  renderMention?: (mention: AIInputMention, token: string) => ReactNode;
  renderVideo?: (src: string, label: string) => ReactNode;
  value: string;
};

type MarkdownOptions = Omit<MarkdownProps, "value">;

/**
 * Splits raw markdown into top-level blocks on blank lines, EXCEPT inside
 * fenced code (``` / ~~~) or display math ($$ ... $$), where blank lines
 * are content, not separators. Naively splitting on /\n\s*\n/ shreds any
 * multi-line fence that contains an internal blank line: the opening fence
 * ends up in one "block" with no closing fence (renders as an empty code
 * block), and the closing fence + remaining code falls into the next
 * block as plain text (renders as literal backticks/JSX in a <p>).
 */
function splitBlocks(value: string): string[] {
  const lines = value.split("\n");
  const blocks: string[] = [];
  let current: string[] = [];
  let mode: "none" | "fence" | "math" = "none";
  let fenceMarker = "";

  const flush = () => {
    const text = current.join("\n").trim();
    if (text) blocks.push(text);
    current = [];
  };

  for (const line of lines) {
    if (mode === "fence") {
      current.push(line);
      if (line.trim() === fenceMarker) {
        mode = "none";
        flush();
      }
      continue;
    }

    if (mode === "math") {
      current.push(line);
      if (line.trim().endsWith("$$")) {
        mode = "none";
        flush();
      }
      continue;
    }

    const fenceMatch = line.match(/^\s*(```|~~~)/);
    if (fenceMatch) {
      flush();
      mode = "fence";
      fenceMarker = fenceMatch[1];
      current.push(line);
      continue;
    }

    const trimmed = line.trim();
    if (trimmed === "$$" || (trimmed.startsWith("$$") && !trimmed.endsWith("$$") === false && trimmed.length === 2)) {
      // handled below for the general case
    }
    if (trimmed.startsWith("$$") && trimmed !== "$$" && trimmed.endsWith("$$") && trimmed.length > 2) {
      // single-line $$...$$ block, closes immediately
      flush();
      blocks.push(line.trim());
      continue;
    }
    if (trimmed === "$$" || (trimmed.startsWith("$$") && !trimmed.endsWith("$$"))) {
      flush();
      mode = "math";
      current.push(line);
      continue;
    }

    if (trimmed === "") {
      flush();
      continue;
    }

    current.push(line);
  }

  flush();
  return blocks;
}

function inline(value: string, props: MarkdownOptions): ReactNode[] {
  const tokens = value.split(/(`[^`]+`|\*\*[^*]+\*\*|\[[^\]]+\]\([^)]+\)|!\[[^\]]*\]\([^)]+\)|[@#/][\w-]+|\$[^$]+\$)/g).filter(Boolean);
  return tokens.map((part, index) => {
    if (part.startsWith("`") && part.endsWith("`")) return <code className="rounded border border-[var(--uhuru-border-subtle)] bg-[var(--uhuru-surface-default)] px-1.5 py-0.5 text-[0.9em] text-[var(--uhuru-text-primary)]" key={index}>{part.slice(1, -1)}</code>;
    if (part.startsWith("**") && part.endsWith("**")) return <strong key={index}>{inline(part.slice(2, -2), props)}</strong>;
    const image = part.match(/^!\[([^\]]*)\]\(([^)]+)\)$/);
    if (image) return props.renderImage?.(image[2], image[1]) ?? <MediaBlock alt={image[1]} key={index} src={image[2]} />;
    const link = part.match(/^\[([^\]]+)\]\(([^)]+)\)$/);
    if (link) {
      const isVideo = /\.(mp4|webm|ogg)(?:\?.*)?$/i.test(link[2]);
      if (isVideo) return props.renderVideo?.(link[2], link[1]) ?? <MediaBlock key={index} src={link[2]} type="video" />;
      return props.renderLink?.(link[2], link[1]) ?? <EmbedBlock href={link[2]} key={index} title={link[1]} />;
    }
    if (part.startsWith("$") && part.endsWith("$")) return props.renderMath?.(part.slice(1, -1), false) ?? <code className="text-[color-mix(in_srgb,var(--uhuru-accent-strong)_92%,transparent)]" key={index}>{part}</code>;
    const mention = part.match(/^([@#/])([\w-]+)$/);
    const matchedMention = mention && props.mentions?.find((item) => (item.trigger ?? mention[1]) === mention[1] && (item.value ?? item.label).toLowerCase() === mention[2].toLowerCase());
    if (mention && matchedMention) {
      return props.renderMention?.(matchedMention, part) ?? (
        <mark
          className={cn("rounded border border-[var(--uhuru-border-accent)] px-1 py-0.5", props.mentionHighlightClassName)}
          key={index}
          style={{
            backgroundColor: "color-mix(in srgb, var(--uhuru-accent-solid) 16%, transparent)",
            color: "var(--uhuru-accent-strong)",
          }}
        >
          {part}
        </mark>
      );
    }
    return part;
  });
}

function renderTable(lines: string[], props: MarkdownOptions, key: string) {
  const cells = (line: string) => line.replace(/^\|/, "").replace(/\|$/, "").split("|").map((cell) => cell.trim());
  const headers = cells(lines[0]);
  const rows = lines.slice(2).map(cells);
  return <TableBlock headers={headers.map((cell) => inline(cell, props))} key={key} rows={rows.map((row) => row.map((cell) => inline(cell, props)))} />;
}

function splitAdvancedLine(line: string) {
  return line.split("|").map((part) => part.trim());
}

function renderAdvancedFence(language: string, source: string, key: string) {
  const lines = source.split("\n").map((line) => line.trim()).filter(Boolean);

  switch (language.toLowerCase()) {
    case "steps":
      return <StepsBlock key={key} steps={lines.map((line) => { const parts = splitAdvancedLine(line.replace(/^\d+[.)]\s*/, "").replace(/^-\s*/, "")); return { description: parts[1], title: parts[0] }; })} />;
    case "checklist":
      return <ChecklistBlock key={key} items={lines.map((line) => ({ checked: /^[-*]?\s*\[[xX]\]/.test(line), label: line.replace(/^[-*]?\s*\[[ xX]\]\s*/, "") }))} />;
    case "key-value":
    case "keyvalue":
      return <KeyValueBlock items={lines.map((line) => { const separator = line.indexOf(":"); return { label: separator === -1 ? line : line.slice(0, separator).trim(), value: separator === -1 ? "" : line.slice(separator + 1).trim() }; })} key={key} />;
    case "timeline":
      return <TimelineBlock items={lines.map((line) => { const [time, title, description] = splitAdvancedLine(line); return { description, time, title }; })} key={key} />;
    case "file": {
      const [name, meta, download] = splitAdvancedLine(lines[0] ?? "");
      return <FileBlock download={download} key={key} meta={meta} name={name} />;
    }
    case "media": {
      const [type, src, alt] = splitAdvancedLine(lines[0] ?? "");
      return <MediaBlock alt={alt} key={key} src={src ?? ""} type={type === "audio" || type === "video" ? type : "image"} />;
    }
    case "embed": {
      const [title, href] = splitAdvancedLine(lines[0] ?? "");
      return <EmbedBlock href={href ?? title} key={key} title={href ? title : undefined} />;
    }
    default:
      return null;
  }
}

export function Markdown({ allowRawHtml = false, className, value, ...props }: MarkdownProps) {
  const blocks = useMemo(() => splitBlocks(value), [value]);

  return <div className={cn("min-w-0 max-w-full space-y-3 break-words text-[var(--uhuru-text-primary)] leading-6 [overflow-wrap:anywhere] [&_img]:max-w-full [&_pre]:max-w-full [&_pre]:overflow-x-auto [&_table]:max-w-full [&_table]:overflow-x-auto", className)}>
    {blocks.map((block, index) => {
      const lines = block.split("\n");
      const fence = lines[0]?.trim().match(/^(```|~~~)\s*([^\s`~]*)/);
      if (fence) {
        const source = lines.slice(1, -1).join("\n");
        const language = fence[2] || "";
        const context = { language: language || undefined, value: source };
        const advanced = renderAdvancedFence(language, source, String(index));
        if (advanced) return advanced;
        if (["mermaid", "diagram", "draw", "drawing"].includes(language.toLowerCase())) return props.renderDiagram?.(source, context) ?? <DiagramBlock key={index} language={language || "diagram"} onCopy={props.onCopy} source={source} />;
        return props.renderCode?.(source, context) ?? <CodeBlock code={source} key={index} language={language} onCopy={props.onCopy} />;
      }
      if (lines[0]?.startsWith("$$") && lines.at(-1)?.endsWith("$$")) { const formula = lines.slice(0, -1).join("\n").replace(/^\$\$/, ""); return props.renderMath?.(formula, true) ?? <MathBlock formula={formula} key={index} onCopy={props.onCopy} />; }
      if (lines.length >= 2 && /^\|.*\|$/.test(lines[0]) && /^\|?\s*:?-+:?\s*(\|\s*:?-+:?\s*)+\|?$/.test(lines[1])) return renderTable(lines, props, String(index));
      if (lines.every((line) => /^[-*] \[[ xX]\] /.test(line))) return <ChecklistBlock items={lines.map((line) => ({ checked: /^[-*] \[[xX]\] /.test(line), label: inline(line.replace(/^[-*] \[[ xX]\] /, ""), props) }))} key={index} />;
      if (lines.every((line) => /^[-*] /.test(line))) return <ul className="list-disc space-y-1 pl-5" key={index}>{lines.map((line) => <li key={line}>{inline(line.slice(2), props)}</li>)}</ul>;
      if (lines.every((line) => /^\d+\. /.test(line))) return <StepsBlock key={index} steps={lines.map((line) => ({ title: inline(line.replace(/^\d+\. /, ""), props) }))} />;
      if (lines.length > 1 && lines.every((line) => /^[^:]+:\s*.+$/.test(line))) return <KeyValueBlock items={lines.map((line) => { const separator = line.indexOf(":"); return { label: inline(line.slice(0, separator), props), value: inline(line.slice(separator + 1).trim(), props) }; })} key={index} />;
      if (lines[0]?.startsWith("> ")) return <QuoteBlock key={index}>{inline(lines.map((line) => line.replace(/^> /, "")).join(" "), props)}</QuoteBlock>;
      const heading = lines[0]?.match(/^(#{1,6})\s+(.*)$/);
      if (heading) { const headingTag = `h${heading[1].length}` as "h1" | "h2" | "h3" | "h4" | "h5" | "h6"; return createElement(headingTag, { className: "font-medium text-[var(--uhuru-text-primary)]", key: index }, inline(heading[2], props)); }
      if (allowRawHtml && /^<([a-z][^>]*)>/i.test(block)) return <div dangerouslySetInnerHTML={{ __html: block }} key={index} />;
      return <p key={index}>{inline(lines.join(" "), props)}</p>;
    })}
  </div>;
}