components/chat/Message.tsx

"use client";

import type { ChatTurn } from "@/lib/chatStore";
import type { SourceDoc } from "@/lib/types";
import { GroundingPill } from "./GroundingPill";
import { SourceChips } from "./SourceChips";
import { useState, type ReactNode } from "react";
import clsx from "clsx";
import ReactMarkdown from "react-markdown";

export function UserMessage({ text }: { text: string }) {
  return (
    <div className="flex justify-end animate-fade-up">
      <div className="max-w-[82%] chat-bubble-user">
        <span className="whitespace-pre-wrap">{text}</span>
      </div>
    </div>
  );
}

export function AssistantMessage({
  turn,
  onOpenSource,
}: {
  turn: ChatTurn;
  onOpenSource?: (doc: SourceDoc, index: number) => void;
}) {
  const r = turn.response;
  const [copied, setCopied] = useState(false);
  if (!r) return null;

  const copy = async () => {
    try {
      await navigator.clipboard.writeText(r.answer);
      setCopied(true);
      setTimeout(() => setCopied(false), 1500);
    } catch {}
  };

  const sources = r.source_docs ?? [];
  const rejected = r._grounding_status === "rejected_low_similarity";

  // Inline citation renderer: replaces [1] [2] markers with clickable badges.
  // Numbers are 1-indexed against `sources` (clamped to length).
  const renderWithCitations = (text: string): ReactNode[] => {
    if (!sources.length || !text) return [text];
    const out: ReactNode[] = [];
    const re = /\[(\d{1,3})\]/g;
    let last = 0;
    let m: RegExpExecArray | null;
    let key = 0;
    while ((m = re.exec(text)) !== null) {
      const idx = parseInt(m[1], 10);
      if (idx >= 1 && idx <= sources.length) {
        if (m.index > last) out.push(text.slice(last, m.index));
        const doc = sources[idx - 1];
        out.push(
          <button
            key={`cite-${key++}`}
            onClick={() => onOpenSource?.(doc, idx - 1)}
            className="inline-flex items-center justify-center align-baseline mx-0.5 px-1.5 min-w-[18px] h-[18px] rounded-full bg-amber/15 text-amber text-[11px] font-mono font-semibold hover:bg-amber hover:text-canvas transition-colors"
            title={`View source: ${doc.name}`}
          >
            {idx}
          </button>
        );
        last = m.index + m[0].length;
      }
    }
    if (last < text.length) out.push(text.slice(last));
    return out.length ? out : [text];
  };

  // Wrap a markdown children array, expanding string nodes through citation regex.
  const transformChildren = (children: ReactNode): ReactNode => {
    if (typeof children === "string") return renderWithCitations(children);
    if (Array.isArray(children))
      return children.map((c, i) =>
        typeof c === "string" ? (
          <span key={i}>{renderWithCitations(c)}</span>
        ) : (
          c
        )
      );
    return children;
  };

  return (
    <div className="flex items-start gap-3 animate-fade-up">
      <Avatar />
      <article className="chat-bubble-assistant flex-1 min-w-0 max-w-[88%]">
        {/* Header: grounding + actions */}
        <header className="flex items-center justify-between gap-3 mb-3">
          <GroundingPill status={r._grounding_status} />
          <div className="flex items-center gap-1">
            <button
              onClick={copy}
              className="text-micro text-ink-50 hover:text-ink hover:bg-glass rounded-md px-2 py-1 transition-colors flex items-center gap-1"
              title="Copy answer"
            >
              {copied ? (
                <>
                  <svg width="11" height="11" viewBox="0 0 11 11" fill="none">
                    <path d="M2 6 4.5 8.5 9 3" stroke="currentColor" strokeWidth="1.6" strokeLinecap="round" strokeLinejoin="round" />
                  </svg>
                  copied
                </>
              ) : (
                <>
                  <svg width="11" height="11" viewBox="0 0 11 11" fill="none">
                    <rect x="2" y="2" width="6" height="6" rx="1" stroke="currentColor" strokeWidth="1.2" />
                    <rect x="3.5" y="3.5" width="6" height="6" rx="1" stroke="currentColor" strokeWidth="1.2" />
                  </svg>
                  copy
                </>
              )}
            </button>
          </div>
        </header>

        {/* Answer */}
        <div className="text-ink leading-[1.65] markdown-body">
          <ReactMarkdown
            components={{
              p: ({ children }) => (
                <p className="mb-3 last:mb-0">{transformChildren(children)}</p>
              ),
              ul: ({ children }) => (
                <ul className="list-disc pl-5 mb-3 last:mb-0 space-y-1 marker:text-amber/70">
                  {children}
                </ul>
              ),
              ol: ({ children }) => (
                <ol className="list-decimal pl-5 mb-3 last:mb-0 space-y-1 marker:text-amber/70 marker:font-mono">
                  {children}
                </ol>
              ),
              li: ({ children }) => <li className="pl-1">{transformChildren(children)}</li>,
              strong: ({ children }) => (
                <strong className="font-semibold text-ink">{children}</strong>
              ),
              em: ({ children }) => (
                <em className="serif-italic text-amber-soft">{children}</em>
              ),
              code: ({ children }) => (
                <code
                  className="font-mono text-amber px-1.5 py-0.5 rounded text-[0.92em]"
                  style={{ background: "rgba(255,181,69,0.10)" }}
                >
                  {children}
                </code>
              ),
              pre: ({ children }) => (
                <pre
                  className="font-mono text-caption p-3 rounded-md overflow-x-auto my-3 leading-[1.55]"
                  style={{
                    background: "rgba(0,0,0,0.28)",
                    boxShadow: "0 0 0 1px rgba(255,255,255,0.06)",
                  }}
                >
                  {children}
                </pre>
              ),
              h1: ({ children }) => (
                <h3 className="text-display-sm mt-4 mb-2">{children}</h3>
              ),
              h2: ({ children }) => (
                <h3 className="text-body-strong text-ink mt-3 mb-1.5">{children}</h3>
              ),
              h3: ({ children }) => (
                <h4 className="text-body-strong text-ink mt-3 mb-1.5">{children}</h4>
              ),
              a: ({ href, children }) => (
                <a
                  href={href}
                  className="text-amber underline decoration-amber/40 underline-offset-2 hover:text-amber-soft hover:decoration-amber transition-colors"
                  target="_blank"
                  rel="noopener noreferrer"
                >
                  {children}
                </a>
              ),
              blockquote: ({ children }) => (
                <blockquote
                  className="border-l-2 border-amber/50 pl-4 my-3 text-ink-70 serif-italic"
                  style={{ fontSize: "0.98em" }}
                >
                  {children}
                </blockquote>
              ),
              hr: () => <hr className="my-4 border-0 h-px bg-glass-border" />,
            }}
          >
            {r.answer}
          </ReactMarkdown>
        </div>

        {/* Sources */}
        {sources.length > 0 && (
          <SourceChips
            docs={sources}
            onOpenSource={onOpenSource}
            rejected={rejected}
            threshold={r._threshold}
          />
        )}

        {/* Spec sheet — collapsed by default */}
        <SpecSheet response={r} />
      </article>
    </div>
  );
}

function SpecSheet({ response: r }: { response: NonNullable<ChatTurn["response"]> }) {
  const [open, setOpen] = useState(false);
  return (
    <footer className="mt-4 pt-3 border-t border-glass-border">
      <button
        onClick={() => setOpen((o) => !o)}
        className="w-full flex items-center justify-between text-micro text-ink-50 font-mono uppercase tracking-[0.10em] hover:text-ink transition-colors"
        aria-expanded={open}
      >
        <span className="flex items-baseline gap-1.5">
          <span>total</span>
          <span className="text-amber">{r.total_seconds.toFixed(2)}s</span>
        </span>
        <span className="flex items-baseline gap-3">
          <span className="hidden sm:inline">
            {open ? "Hide diagnostics" : "Show diagnostics"}
          </span>
          <span
            className={clsx(
              "transition-transform duration-200 inline-block",
              open && "rotate-180"
            )}
            aria-hidden
          >
            <svg width="9" height="9" viewBox="0 0 9 9" fill="none">
              <path
                d="m2 3 2.5 2.5L7 3"
                stroke="currentColor"
                strokeWidth="1.4"
                strokeLinecap="round"
              />
            </svg>
          </span>
        </span>
      </button>
      {open && (
        <div className="mt-3 grid grid-cols-3 sm:grid-cols-5 gap-x-5 gap-y-2 text-micro text-ink-50 font-mono uppercase tracking-[0.10em] animate-fade-in">
          <Spec label="top·sim" value={r._top_similarity?.toFixed(3) ?? "—"} />
          <Spec label="rerank" value={r._top_rerank_score?.toFixed(2) ?? "—"} />
          <Spec label="anchor" value={r._anchor_score?.toFixed(3) ?? "—"} />
          <Spec label="retrieve" value={`${r.retrieve_seconds.toFixed(2)}s`} />
          <Spec label="inference" value={`${r.inference_seconds.toFixed(2)}s`} />
        </div>
      )}
    </footer>
  );
}

export function ErrorMessage({ turn }: { turn: ChatTurn }) {
  const e = turn.error;
  if (!e) return null;
  return (
    <div className="flex items-start gap-3 animate-fade-up">
      <Avatar tone="err" />
      <article
        className="flex-1 min-w-0 max-w-[88%] rounded-2xl rounded-tl-sm px-5 py-4"
        style={{
          background: "rgba(255, 122, 122, 0.06)",
          boxShadow: "0 0 0 1px rgba(255, 122, 122, 0.28)",
        }}
      >
        <div className="text-micro uppercase tracking-[0.12em] text-status-err font-mono">
          error · http {e.status || "—"}
        </div>
        <div className="text-body-strong text-ink mt-1.5">{e.message}</div>
        {e.status === 502 && (
          <div className="text-caption text-ink-70 mt-2">
            Backend not reachable. Check the HF Space stage.
          </div>
        )}
      </article>
    </div>
  );
}

function Spec({ label, value }: { label: string; value: string }) {
  return (
    <div className="flex flex-col">
      <span className="text-ink-50">{label}</span>
      <span className="text-ink mt-0.5 normal-case tracking-normal">{value}</span>
    </div>
  );
}

function Avatar({ tone = "amber" }: { tone?: "amber" | "err" }) {
  const bg =
    tone === "err"
      ? "linear-gradient(135deg, rgba(255,122,122,0.95), rgba(195,70,70,1))"
      : "linear-gradient(135deg, rgba(255,181,69,0.95), rgba(214,138,31,1))";
  const ring =
    tone === "err"
      ? "0 0 0 1px rgba(255,122,122,0.45), 0 0 14px -2px rgba(255,122,122,0.4)"
      : "0 0 0 1px rgba(255,181,69,0.45), 0 0 14px -2px rgba(255,181,69,0.45)";
  return (
    <div
      className="shrink-0 w-7 h-7 rounded-md flex items-center justify-center mt-1"
      style={{ background: bg, boxShadow: ring }}
    >
      <svg width="14" height="14" viewBox="0 0 14 14" fill="none">
        <path
          d="M2.5 7.5 4.8 9.8l6-6.6"
          stroke="#0a0b10"
          strokeWidth="1.8"
          strokeLinecap="round"
          strokeLinejoin="round"
        />
      </svg>
    </div>
  );
}