freezmo/components/chat/Markdown.tsx

"use client";

import { memo, useMemo } from "react";

import { CopyIconButton } from "./CopyIconButton";
import { MathSpan } from "./MathSpan";

/**
 * Safe, dependency-free markdown for chat answers: lists, tables, code blocks
 * (with a copy button), quotes, headings, inline code/bold/italic/links and
 * $...$ math. It parses to React nodes — nothing is ever set as HTML.
 */

type Block =
  | { type: "p"; lines: string[] }
  | { type: "heading"; level: number; text: string }
  | { type: "list"; ordered: boolean; items: { text: string; depth: number }[] }
  | { type: "code"; lang: string; code: string }
  | { type: "table"; head: string[]; rows: string[][] }
  | { type: "quote"; lines: string[] }
  | { type: "math"; tex: string }
  | { type: "hr" };

const FENCE = /^(```|~~~)\s*([\w+-]*)\s*$/;
const HEADING = /^(#{1,6})\s+(.*)$/;
const LIST_ITEM = /^(\s*)([-*+]|\d+[.)])\s+(.*)$/;
const QUOTE = /^>\s?(.*)$/;
const HR = /^\s*([-*_])\1{2,}\s*$/;

function isTableSeparator(line: string): boolean {
  return /^\s*\|?[\s:|-]*-[\s:|-]*\|?\s*$/.test(line) && line.includes("-");
}

function splitTableRow(line: string): string[] {
  const trimmed = line.trim().replace(/^\|/, "").replace(/\|$/, "");
  return trimmed.split("|").map((cell) => cell.trim());
}

export function parseBlocks(text: string): Block[] {
  const blocks: Block[] = [];
  const lines = text.replace(/\r\n/g, "\n").split("\n");
  let index = 0;

  const flushParagraph = (collected: string[]) => {
    const joined = collected.join("\n").trim();
    if (joined) blocks.push({ type: "p", lines: collected });
  };

  let paragraph: string[] = [];

  while (index < lines.length) {
    const line = lines[index];

    const fence = FENCE.exec(line);
    if (fence) {
      flushParagraph(paragraph);
      paragraph = [];
      const lang = fence[2] ?? "";
      const code: string[] = [];
      index += 1;
      while (index < lines.length && !FENCE.test(lines[index])) {
        code.push(lines[index]);
        index += 1;
      }
      index += 1; // closing fence (or end of stream)
      blocks.push({ type: "code", lang, code: code.join("\n") });
      continue;
    }

    if (HR.test(line)) {
      flushParagraph(paragraph);
      paragraph = [];
      blocks.push({ type: "hr" });
      index += 1;
      continue;
    }

    const heading = HEADING.exec(line);
    if (heading) {
      flushParagraph(paragraph);
      paragraph = [];
      blocks.push({ type: "heading", level: heading[1].length, text: heading[2].trim() });
      index += 1;
      continue;
    }

    // $$ ... $$ on its own line(s)
    if (line.trim().startsWith("$$")) {
      flushParagraph(paragraph);
      paragraph = [];
      const trimmed = line.trim();
      if (trimmed.length > 4 && trimmed.endsWith("$$")) {
        blocks.push({ type: "math", tex: trimmed.slice(2, -2).trim() });
        index += 1;
        continue;
      }
      const math: string[] = [trimmed.slice(2)];
      index += 1;
      while (index < lines.length && !lines[index].trim().endsWith("$$")) {
        math.push(lines[index]);
        index += 1;
      }
      if (index < lines.length) {
        math.push(lines[index].trim().slice(0, -2));
        index += 1;
      }
      blocks.push({ type: "math", tex: math.join("\n").trim() });
      continue;
    }

    const quote = QUOTE.exec(line);
    if (quote) {
      flushParagraph(paragraph);
      paragraph = [];
      const collected = [quote[1]];
      index += 1;
      while (index < lines.length) {
        const next = QUOTE.exec(lines[index]);
        if (!next) break;
        collected.push(next[1]);
        index += 1;
      }
      blocks.push({ type: "quote", lines: collected });
      continue;
    }

    const item = LIST_ITEM.exec(line);
    if (item) {
      flushParagraph(paragraph);
      paragraph = [];
      const ordered = /\d/.test(item[2]);
      const items: { text: string; depth: number }[] = [];
      while (index < lines.length) {
        const next = LIST_ITEM.exec(lines[index]);
        if (!next) break;
        const depth = Math.min(2, Math.floor(next[1].replace(/\t/g, "  ").length / 2));
        items.push({ text: next[3], depth });
        index += 1;
      }
      blocks.push({ type: "list", ordered, items });
      continue;
    }

    // Table: header row followed by a separator row.
    if (line.includes("|") && index + 1 < lines.length && isTableSeparator(lines[index + 1])) {
      flushParagraph(paragraph);
      paragraph = [];
      const head = splitTableRow(line);
      index += 2;
      const rows: string[][] = [];
      while (index < lines.length && lines[index].includes("|") && lines[index].trim()) {
        rows.push(splitTableRow(lines[index]));
        index += 1;
      }
      blocks.push({ type: "table", head, rows });
      continue;
    }

    if (!line.trim()) {
      flushParagraph(paragraph);
      paragraph = [];
      index += 1;
      continue;
    }

    paragraph.push(line);
    index += 1;
  }
  flushParagraph(paragraph);
  return blocks;
}

const INLINE =
  /(\*\*[^*]+\*\*|__[^_]+__|`[^`]+`|\[[^\]]+\]\([^)\s]+\)|\$\$[^$]+\$\$|\$[^$\n]+\$|\*[^*\n]+\*|_[^_\n]+_)/g;

function inlineNodes(text: string, keyPrefix: string): React.ReactNode[] {
  return text.split(INLINE).map((part, index) => {
    const key = `${keyPrefix}-${index}`;
    if (!part) return null;
    if (part.startsWith("**") && part.endsWith("**") && part.length > 4) {
      return <strong key={key}>{part.slice(2, -2)}</strong>;
    }
    if (part.startsWith("__") && part.endsWith("__") && part.length > 4) {
      return <strong key={key}>{part.slice(2, -2)}</strong>;
    }
    if (part.startsWith("`") && part.endsWith("`") && part.length > 2) {
      return (
        <code key={key} className="rounded-md bg-[var(--surface-sunken)] px-1.5 py-0.5 font-mono text-[0.85em]">
          {part.slice(1, -1)}
        </code>
      );
    }
    if (part.startsWith("$$") && part.endsWith("$$") && part.length > 4) {
      return <MathSpan key={key} tex={part.slice(2, -2).trim()} display />;
    }
    if (part.startsWith("$") && part.endsWith("$") && part.length > 2) {
      return <MathSpan key={key} tex={part.slice(1, -1)} />;
    }
    const link = /^\[([^\]]+)\]\(([^)\s]+)\)$/.exec(part);
    if (link) {
      const href = link[2];
      if (/^https?:\/\//i.test(href)) {
        return (
          <a
            key={key}
            href={href}
            target="_blank"
            rel="noopener noreferrer"
            className="font-medium text-accent-600 underline decoration-accent-400/50 underline-offset-2 dark:text-accent-300"
          >
            {link[1]}
          </a>
        );
      }
      return <span key={key}>{link[1]}</span>;
    }
    if (part.startsWith("*") && part.endsWith("*") && part.length > 2) {
      return <em key={key}>{part.slice(1, -1)}</em>;
    }
    if (part.startsWith("_") && part.endsWith("_") && part.length > 2) {
      return <em key={key}>{part.slice(1, -1)}</em>;
    }
    return <span key={key}>{part}</span>;
  });
}

function renderBlock(block: Block, index: number) {
  const key = `b-${index}`;
  switch (block.type) {
    case "heading": {
      const Tag = (block.level <= 2 ? "h3" : "h4") as "h3" | "h4";
      return (
        <Tag key={key} className="text-[15px] font-semibold text-[var(--text)]">
          {inlineNodes(block.text, key)}
        </Tag>
      );
    }
    case "list": {
      const Tag = block.ordered ? "ol" : "ul";
      return (
        <Tag key={key} className={block.ordered ? "list-decimal space-y-1 pl-5" : "list-disc space-y-1 pl-5"}>
          {block.items.map((item, itemIndex) => (
            <li
              key={`${key}-${itemIndex}`}
              className={item.depth > 0 ? "ml-4 list-[circle]" : undefined}
            >
              {inlineNodes(item.text, `${key}-${itemIndex}`)}
            </li>
          ))}
        </Tag>
      );
    }
    case "code":
      return (
        <div key={key} className="group relative">
          <pre className="overflow-x-auto rounded-xl border border-[var(--hairline-soft)] bg-[var(--surface-sunken)] p-3 text-[12.5px] leading-relaxed">
            <code className="font-mono">{block.code}</code>
          </pre>
          <div className="absolute right-2 top-2 opacity-90">
            <CopyIconButton text={block.code} />
          </div>
        </div>
      );
    case "table":
      return (
        <div key={key} className="overflow-x-auto">
          <table className="w-full min-w-[280px] border-collapse text-[13px]">
            <thead>
              <tr>
                {block.head.map((cell, cellIndex) => (
                  <th
                    key={`${key}-h-${cellIndex}`}
                    className="border-b border-[var(--hairline)] px-2 py-1.5 text-left font-semibold text-[var(--text)]"
                  >
                    {inlineNodes(cell, `${key}-h-${cellIndex}`)}
                  </th>
                ))}
              </tr>
            </thead>
            <tbody>
              {block.rows.map((row, rowIndex) => (
                <tr key={`${key}-r-${rowIndex}`}>
                  {block.head.map((_, cellIndex) => (
                    <td key={`${key}-r-${rowIndex}-${cellIndex}`} className="border-b border-[var(--hairline-soft)] px-2 py-1.5 align-top">
                      {inlineNodes(row[cellIndex] ?? "", `${key}-r-${rowIndex}-${cellIndex}`)}
                    </td>
                  ))}
                </tr>
              ))}
            </tbody>
          </table>
        </div>
      );
    case "quote":
      return (
        <blockquote key={key} className="border-l-2 border-[var(--hairline)] pl-3 text-[var(--text-muted)]">
          {inlineNodes(block.lines.join("\n"), key)}
        </blockquote>
      );
    case "math":
      return <MathSpan key={key} tex={block.tex} display />;
    case "hr":
      return <hr key={key} className="border-[var(--hairline-soft)]" />;
    default:
      return (
        <p key={key} className="whitespace-pre-wrap break-words">
          {inlineNodes(block.lines.join("\n"), key)}
        </p>
      );
  }
}

export const Markdown = memo(function Markdown({ text }: { text: string }) {
  const blocks = useMemo(() => parseBlocks(text), [text]);
  return <div className="chat-md space-y-3 text-[15px] leading-relaxed text-[var(--text)]">{blocks.map(renderBlock)}</div>;
});