import * as React from "react";
import ReactMarkdown, { type Components } from "react-markdown";
import remarkGfm from "remark-gfm";
import rehypeSanitize from "rehype-sanitize";
import { cn } from "@/lib/utils";

/**
 * MarkdownContent — WealthX DS (Atom)
 *
 * Renders a markdown string into safe, design-system-styled HTML. Used to render
 * AI / assistant responses that arrive as markdown (headings, lists, bold, links,
 * tables, code blocks) instead of plain text.
 *
 * Safety: every render runs through `rehype-sanitize` with the default schema, so
 * untrusted model output cannot inject scripts, event handlers, or unsafe URLs.
 *
 * Styling is applied per-element via the `components` map (the repo does not ship
 * `@tailwindcss/typography`), using design tokens only and sharp corners.
 */

export interface MarkdownContentProps {
  /** The markdown source string to render. */
  content: string;
  /** Density of the rendered text. Defaults to `"medium"`. */
  size?: "small" | "medium";
  /**
   * Known mention labels (e.g. staff names). Any `@<label>` run in the text is
   * rendered as a primary-coloured tag. Matched against this list so labels
   * containing spaces (e.g. "@Iris Chen") have a known boundary.
   *
   * Only needed for plain-text mentions. Identity-carrying mentions written as
   * `[@Name](#staff-id)` by `ChatInputArea` render as chips without this prop —
   * prefer those: a name match cannot tell two people apart or survive a rename.
   */
  mentions?: string[];
  /**
   * Render `[@Name](#staff-id)` links (as written by `ChatInputArea`) as mention chips
   * instead of links. Defaults to `false` so a surface that cannot author mentions — AI
   * replies, assistant output — never has a stray `#staff-` link stop behaving like one.
   *
   * Turn this on wherever people tag each other, e.g. a comments thread.
   */
  renderMentions?: boolean;
  className?: string;
}

const escapeRegExp = (s: string): string =>
  s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");

/** Href prefix marking a link as a mention rather than a destination. */
const MENTION_HREF_PREFIX = "#staff-";

/** Wraps `@<label>` runs (for known labels) in a primary-coloured tag. */
function highlightMentions(
  children: React.ReactNode,
  splitRe: RegExp,
  testRe: RegExp,
): React.ReactNode {
  return React.Children.map(children, (child) => {
    if (typeof child !== "string") return child;
    return child.split(splitRe).map((part, i) =>
      testRe.test(part) ? (
        <span key={i} className="font-medium text-primary">
          {part}
        </span>
      ) : (
        part
      ),
    );
  });
}

// Per-element renderers. Hoisted to module scope so react-markdown receives a
// stable `components` reference (this atom re-renders on every streamed token).
const MARKDOWN_COMPONENTS: Components = {
  p: ({ node, ...props }) => <p className="m-0" {...props} />,
  h1: ({ node, ...props }) => (
    <h1 className="text-body-large font-semibold" {...props} />
  ),
  h2: ({ node, ...props }) => (
    <h2 className="text-body-large font-semibold" {...props} />
  ),
  h3: ({ node, ...props }) => (
    <h3 className="text-body-medium font-semibold" {...props} />
  ),
  ul: ({ node, ...props }) => (
    <ul className="m-0 flex list-disc flex-col gap-1 pl-5" {...props} />
  ),
  ol: ({ node, ...props }) => (
    <ol className="m-0 flex list-decimal flex-col gap-1 pl-5" {...props} />
  ),
  li: ({ node, ...props }) => <li className="pl-0.5" {...props} />,
  a: ({ node, ...props }) => (
    <a
      className="font-medium text-primary underline underline-offset-2 hover:text-primary/80"
      target="_blank"
      rel="noopener noreferrer"
      {...props}
    />
  ),
  strong: ({ node, ...props }) => (
    <strong className="font-semibold" {...props} />
  ),
  em: ({ node, ...props }) => <em className="italic" {...props} />,
  blockquote: ({ node, ...props }) => (
    <blockquote
      className="border-l-2 border-border pl-3 text-muted-foreground"
      {...props}
    />
  ),
  code: ({ node, className: codeClass, children, ...props }) => {
    const isBlock = codeClass?.includes("language-");
    if (isBlock) {
      return (
        <code
          className={cn("block font-mono text-body-small", codeClass)}
          {...props}
        >
          {children}
        </code>
      );
    }
    return (
      <code className="bg-muted px-1 py-0.5 font-mono text-[0.85em]" {...props}>
        {children}
      </code>
    );
  },
  pre: ({ node, ...props }) => (
    <pre
      className="overflow-x-auto border border-border bg-muted p-3 font-mono text-body-small"
      {...props}
    />
  ),
  table: ({ node, ...props }) => (
    <div className="overflow-x-auto">
      <table
        className="w-full border-collapse border border-border text-body-small"
        {...props}
      />
    </div>
  ),
  thead: ({ node, ...props }) => <thead className="bg-muted" {...props} />,
  th: ({ node, ...props }) => (
    <th
      className="border border-border px-2 py-1 text-left font-semibold"
      {...props}
    />
  ),
  td: ({ node, ...props }) => (
    <td className="border border-border px-2 py-1 align-top" {...props} />
  ),
  hr: ({ node, ...props }) => (
    <hr className="border-t border-border" {...props} />
  ),
};

/**
 * Opt-in `a` renderer that turns `[@Name](#staff-id)` into a mention chip.
 *
 * Kept behind `renderMentions` rather than always on: this atom also renders untrusted
 * model output, and a surface that cannot author mentions should not have a stray
 * `#staff-` link silently stop being a link.
 */
const MENTION_LINK_COMPONENTS: Components = {
  a: ({ node, href, children, ...props }) => {
    if (href?.startsWith(MENTION_HREF_PREFIX)) {
      // A mention names a person, not somewhere to navigate — drop the href entirely.
      return (
        <span
          className="font-medium text-primary"
          data-mention-id={href.slice(MENTION_HREF_PREFIX.length)}
        >
          {children}
        </span>
      );
    }
    return (
      <a
        className="font-medium text-primary underline underline-offset-2 hover:text-primary/80"
        href={href}
        rel="noopener noreferrer"
        target="_blank"
        {...props}
      >
        {children}
      </a>
    );
  },
};

export function MarkdownContent({
  content,
  size = "medium",
  mentions,
  renderMentions = false,
  className,
}: MarkdownContentProps) {
  const textSize = size === "small" ? "text-body-small" : "text-body-medium";

  // Two independent mention concerns compose here: `renderMentions` swaps the link
  // renderer for the chip-aware one, and `mentions` wraps the text-bearing renderers so
  // plain `@<label>` runs highlight. Neither in play → the stable static map.
  const components = React.useMemo<Components>(() => {
    const base: Components = renderMentions
      ? { ...MARKDOWN_COMPONENTS, ...MENTION_LINK_COMPONENTS }
      : MARKDOWN_COMPONENTS;
    if (!mentions || mentions.length === 0) return base;
    const alts = mentions
      .slice()
      .sort((a, b) => b.length - a.length)
      .map(escapeRegExp)
      .join("|");
    const splitRe = new RegExp(`(@(?:${alts}))`, "g");
    const testRe = new RegExp(`^@(?:${alts})$`);
    return {
      ...base,
      p: ({ node, children, ...props }) => (
        <p className="m-0" {...props}>
          {highlightMentions(children, splitRe, testRe)}
        </p>
      ),
      li: ({ node, children, ...props }) => (
        <li className="pl-0.5" {...props}>
          {highlightMentions(children, splitRe, testRe)}
        </li>
      ),
      td: ({ node, children, ...props }) => (
        <td className="border border-border px-2 py-1 align-top" {...props}>
          {highlightMentions(children, splitRe, testRe)}
        </td>
      ),
    };
  }, [mentions, renderMentions]);

  return (
    <div
      className={cn(
        "flex flex-col gap-2 break-words leading-relaxed",
        textSize,
        className,
      )}
    >
      <ReactMarkdown
        remarkPlugins={[remarkGfm]}
        rehypePlugins={[rehypeSanitize]}
        components={components}
      >
        {content}
      </ReactMarkdown>
    </div>
  );
}