RAG retrievalNew

"use client";

import { Fragment, useEffect, useId, useRef, useState } from "react";
import { Check, FileText, LoaderCircle, RefreshCw, ScanSearch } from "lucide-react";
import { AnimatePresence, motion, useReducedMotion } from "motion/react";
import { cn } from "@/lib/utils";

export type RagChunk = {
  id: string;
  /** Similarity score, 0–1. */
  score: number;
  source: string;
  /** Section, heading or page, e.g. "p.2". */
  locator?: string;
  /** Snippet. Wrap matched terms in **double asterisks** to highlight them. */
  text: string;
  tokens?: number;
  stale?: boolean;
};

export type RagRetrievalProps = {
  query?: string;
  chunks?: RagChunk[];
  /** Chunks start excluded. */
  defaultExcluded?: string[];
  /** Score under which a chunk is flagged as weak. */
  minScore?: number;
  searchedCount?: number;
  onToggle?: (id: string, included: boolean) => void;
  onRerun?: (includedIds: string[]) => void;
  /** Loop a scripted search → inspect → exclude a stale chunk → re-run demo. */
  autoPlay?: boolean;
  className?: string;
};

const DEFAULT_CHUNKS: RagChunk[] = [
  { id: "k1", score: 0.82, source: "billing/refunds.md", locator: "§ Duplicate charges", text: "If a **duplicate charge** lands in the same cycle, the second capture is voided within 24h.", tokens: 210 },
  { id: "k2", score: 0.77, source: "incidents/2026-10-14.md", locator: "Webhook retries", text: "Payment webhook retried 3× on Oct 14, creating a **duplicate charge** for 41 accounts in **October**.", tokens: 380 },
  { id: "k3", score: 0.71, source: "INV-2084.pdf", locator: "p.2", text: "Pro plan · Oct 1–31 · $480.00 — captured twice, second capture pending void.", tokens: 160 },
  { id: "k4", score: 0.58, source: "support/macros.md", locator: "Refund macro", text: "Apologize, confirm the **duplicate charge**, then issue the refund from the dashboard.", tokens: 140 },
  { id: "k5", score: 0.41, source: "faq/pricing-2023.md", locator: "Legacy billing", text: "Charges in **October** 2023 ran on the legacy invoicing system and were batched weekly.", tokens: 290, stale: true },
];

const EASE = [0.22, 1, 0.36, 1] as const;
const TICK = 100;
const LOOP = 10400;

function useDemoClock(enabled: boolean, loop: number) {
  const ref = useRef<HTMLDivElement>(null);
  const [visible, setVisible] = useState(false);
  const [elapsed, setElapsed] = useState(0);
  useEffect(() => {
    const el = ref.current;
    if (!el || !enabled) return;
    const io = new IntersectionObserver(([entry]) => setVisible(entry.isIntersecting));
    io.observe(el);
    return () => io.disconnect();
  }, [enabled]);
  useEffect(() => {
    if (!enabled || !visible) return;
    const id = setInterval(() => setElapsed((e) => e + TICK), TICK);
    return () => clearInterval(id);
  }, [enabled, visible]);
  return { ref, t: elapsed % loop };
}

type Phase = "searching" | "results" | "rerunning" | "updated";
type View = { phase: Phase; stage: 0 | 1; excluded: string[]; focus: string | null; press: boolean };

function demoView(t: number): View {
  if (t < 1700) return { phase: "searching", stage: t < 800 ? 0 : 1, excluded: [], focus: null, press: false };
  const excluded = t >= 4300 ? ["k5"] : [];
  const phase: Phase = t >= 6700 ? "updated" : t >= 5400 ? "rerunning" : "results";
  return { phase, stage: 1, excluded, focus: t >= 3500 && t < 5000 ? "k5" : null, press: (t >= 4100 && t < 4300) || (t >= 5200 && t < 5400) };
}

function Highlight({ text }: { text: string }) {
  const parts = text.split(/\*\*(.+?)\*\*/g);
  return (
    <>
      {parts.map((p, i) =>
        i % 2 ? (
          <mark key={i} className="rounded-[3px] bg-amber-300/40 px-0.5 text-foreground dark:bg-amber-400/25">
            {p}
          </mark>
        ) : (
          <Fragment key={i}>{p}</Fragment>
        ),
      )}
    </>
  );
}

/** Retrieved chunks for a RAG answer: similarity scores, sources, highlighted matches, and per-chunk toggles to re-run the answer. */
export function RagRetrieval({
  query = "duplicate charge october",
  chunks = DEFAULT_CHUNKS,
  defaultExcluded = [],
  minScore = 0.5,
  searchedCount = 2914,
  onToggle,
  onRerun,
  autoPlay = false,
  className,
}: RagRetrievalProps) {
  const reduce = useReducedMotion() ?? false;
  const uid = useId();
  const [touched, setTouched] = useState(false);
  const demo = autoPlay && !reduce && !touched && chunks === DEFAULT_CHUNKS;
  const { ref, t } = useDemoClock(demo, LOOP);
  const [view, setView] = useState<View>({ phase: "results", stage: 1, excluded: defaultExcluded, focus: null, press: false });
  // Sync the scripted frame during render (no effect → no extra commit per tick).
  const [syncedT, setSyncedT] = useState(-1);
  if (demo && t !== syncedT) {
    setSyncedT(t);
    setView(demoView(t));
  }

  const { phase, stage, excluded, focus, press } = view;

  // Interactive re-run settles after a short beat.
  useEffect(() => {
    if (demo || phase !== "rerunning") return;
    const id = setTimeout(() => setView((v) => ({ ...v, phase: "updated" })), reduce ? 0 : 1200);
    return () => clearTimeout(id);
  }, [demo, phase, reduce]);

  const sorted = [...chunks].sort((a, b) => b.score - a.score);
  const included = sorted.filter((c) => !excluded.includes(c.id));
  const tok = included.reduce((s, c) => s + (c.tokens ?? 0), 0);
  const cited = included.slice(0, 3).map((c) => sorted.indexOf(c) + 1);

  const toggle = (id: string) => {
    setTouched(true);
    const on = excluded.includes(id);
    onToggle?.(id, on);
    setView((v) => ({ ...v, excluded: on ? v.excluded.filter((x) => x !== id) : [...v.excluded, id], phase: v.phase === "updated" ? "results" : v.phase, focus: null }));
  };

  const rerun = () => {
    setTouched(true);
    onRerun?.(included.map((c) => c.id));
    setView((v) => ({ ...v, phase: "rerunning", focus: null }));
  };

  return (
    <div ref={ref} className={cn("w-full max-w-md text-sm", className)}>
      <div className="overflow-hidden rounded-2xl border bg-card text-card-foreground shadow-[0_1px_2px_rgb(0_0_0/0.04),0_12px_32px_-18px_rgb(0_0_0/0.25)]">
        <div className="flex items-center gap-3 bg-gradient-to-b from-muted/50 to-transparent px-4 pt-4 pb-3">
          <span className="relative grid size-9 shrink-0 place-items-center overflow-hidden rounded-xl border bg-background shadow-sm">
            <ScanSearch className="relative size-4" aria-hidden />
            {phase === "searching" && !reduce && <span className="absolute inset-x-0 h-3 animate-[rag-scan_1.1s_ease-in-out_infinite] bg-gradient-to-b from-transparent via-sky-500/30 to-transparent" aria-hidden />}
          </span>
          <div className="min-w-0 flex-1">
            <p className="text-sm font-semibold">Retrieved context</p>
            <p className="truncate font-mono text-[10.5px] text-muted-foreground">q=“{query}”</p>
          </div>
          <span className="shrink-0 rounded-full border bg-background px-2 py-0.5 font-mono text-[10px] text-muted-foreground tabular-nums" aria-live="polite">
            {phase === "searching" ? (stage === 0 ? "embedding…" : `scanning ${searchedCount.toLocaleString()}`) : `${included.length}/${chunks.length} · ${(tok / 1000).toFixed(1)}k tok`}
          </span>
        </div>

        <ol className="px-2 pb-1" aria-label="Retrieved chunks">
          {phase === "searching"
            ? sorted.map((c, i) => (
                <li key={c.id} className="flex h-12 items-center gap-3 px-2" aria-hidden>
                  <span className="h-6 w-6 shrink-0 rounded-md bg-muted" />
                  <span className="flex-1 space-y-1.5">
                    <span className="relative block h-2 w-2/5 overflow-hidden rounded-full bg-muted">
                      {!reduce && <span className="absolute inset-0 animate-[rag-shim_1.2s_ease-in-out_infinite] bg-gradient-to-r from-transparent via-foreground/10 to-transparent" style={{ animationDelay: `${i * 0.1}s` }} />}
                    </span>
                    <span className="relative block h-2 w-4/5 overflow-hidden rounded-full bg-muted">
                      {!reduce && <span className="absolute inset-0 animate-[rag-shim_1.2s_ease-in-out_infinite] bg-gradient-to-r from-transparent via-foreground/10 to-transparent" style={{ animationDelay: `${i * 0.1 + 0.05}s` }} />}
                    </span>
                  </span>
                </li>
              ))
            : sorted.map((c, i) => {
                const off = excluded.includes(c.id);
                const weak = c.score < minScore;
                const tone = c.score >= 0.7 ? "from-emerald-500 to-teal-400" : c.score >= minScore ? "from-amber-500 to-yellow-400" : "from-rose-500 to-orange-400";
                return (
                  <motion.li
                    key={c.id}
                    initial={reduce ? false : { opacity: 0, y: 8, filter: "blur(4px)" }}
                    animate={{ opacity: off ? 0.45 : 1, y: 0, filter: "blur(0px)" }}
                    transition={{ delay: phase === "results" && !touched && excluded.length === 0 ? i * 0.08 : 0, duration: 0.4, ease: EASE }}
                    className={cn("relative flex h-12 items-center gap-3 rounded-xl px-2 transition-colors", focus === c.id ? "bg-rose-500/5 ring-1 ring-rose-500/30" : "hover:bg-muted/40")}
                  >
                    <span className="relative grid size-6 shrink-0 place-items-center rounded-md border bg-background font-mono text-[10px] font-semibold tabular-nums">{i + 1}</span>
                    <div className="min-w-0 flex-1">
                      <div className="flex items-center gap-1.5 text-[10.5px]">
                        <FileText className="size-3 shrink-0 text-muted-foreground" aria-hidden />
                        <span className="truncate font-mono">{c.source}</span>
                        {c.locator && <span className="shrink-0 truncate text-muted-foreground">· {c.locator}</span>}
                        {c.stale && <span className="shrink-0 rounded-full bg-rose-500/10 px-1.5 text-[9px] font-medium text-rose-700 dark:text-rose-300">stale</span>}
                      </div>
                      <p className={cn("mt-0.5 truncate text-[11px] text-muted-foreground", off && "line-through decoration-foreground/30")}>
                        <Highlight text={c.text} />
                      </p>
                    </div>
                    <div className="flex w-14 shrink-0 flex-col items-end gap-1">
                      <span className={cn("font-mono text-[10.5px] font-semibold tabular-nums", weak && "text-rose-600 dark:text-rose-400")}>{c.score.toFixed(2)}</span>
                      <span className="relative h-1 w-full overflow-hidden rounded-full bg-muted">
                        <motion.span
                          className={cn("absolute inset-y-0 left-0 rounded-full bg-gradient-to-r", tone)}
                          initial={reduce ? false : { width: 0 }}
                          animate={{ width: `${c.score * 100}%` }}
                          transition={{ delay: 0.15 + i * 0.08, duration: 0.6, ease: EASE }}
                        />
                        <span className="absolute inset-y-0 w-px bg-foreground/40" style={{ left: `${minScore * 100}%` }} aria-hidden />
                      </span>
                    </div>
                    <button
                      type="button"
                      role="switch"
                      aria-checked={!off}
                      aria-label={`Use chunk ${i + 1} from ${c.source}`}
                      onClick={() => toggle(c.id)}
                      className={cn(
                        "relative h-[18px] w-8 shrink-0 cursor-pointer rounded-full outline-none transition-colors duration-300 focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2",
                        off ? "bg-muted-foreground/25" : "bg-foreground",
                        press && focus === c.id && "scale-90",
                      )}
                    >
                      <motion.span
                        className="absolute top-[2px] left-[2px] size-[14px] rounded-full bg-background shadow-sm"
                        initial={false}
                        animate={{ x: off ? 0 : 14 }}
                        transition={{ type: "spring", stiffness: 600, damping: 32 }}
                      />
                    </button>
                  </motion.li>
                );
              })}
        </ol>

        <div className="flex h-12 items-center gap-2 border-t px-3">
          <AnimatePresence mode="popLayout" initial={false}>
            <motion.p
              key={phase === "updated" ? "u" : phase === "rerunning" ? "r" : "n"}
              initial={reduce ? false : { opacity: 0, y: 6 }}
              animate={{ opacity: 1, y: 0 }}
              exit={reduce ? undefined : { opacity: 0, y: -6 }}
              transition={{ duration: 0.25, ease: EASE }}
              className="flex min-w-0 flex-1 items-center gap-1.5 text-[11px]"
              role="status"
            >
              {phase === "updated" ? (
                <>
                  <span className="grid size-4 shrink-0 place-items-center rounded-full bg-emerald-500 text-white">
                    <Check className="size-2.5" strokeWidth={3.5} aria-hidden />
                  </span>
                  <span className="truncate">
                    Answer updated · cites{" "}
                    {cited.map((n) => (
                      <span key={n} className="mr-0.5 inline-grid size-4 place-items-center rounded bg-muted font-mono text-[9px] font-semibold">
                        {n}
                      </span>
                    ))}
                  </span>
                </>
              ) : phase === "rerunning" ? (
                <span className="text-muted-foreground">Regenerating with {included.length} chunks…</span>
              ) : phase === "searching" ? (
                <span className="text-muted-foreground">Ranking by cosine similarity…</span>
              ) : (
                <span className="truncate text-muted-foreground">
                  {excluded.length ? `${excluded.length} chunk excluded · ` : ""}min score <span className="font-mono">{minScore.toFixed(2)}</span>
                </span>
              )}
            </motion.p>
          </AnimatePresence>
          <motion.button
            type="button"
            onClick={rerun}
            disabled={phase === "searching" || phase === "rerunning"}
            animate={{ scale: press && !focus ? 0.94 : 1 }}
            transition={{ type: "spring", stiffness: 500, damping: 26 }}
            aria-describedby={`${uid}-hint`}
            className="inline-flex h-8 shrink-0 cursor-pointer items-center gap-1.5 rounded-lg bg-foreground px-3 text-xs font-medium text-background outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50"
          >
            {phase === "rerunning" ? <LoaderCircle className={cn("size-3.5", !reduce && "animate-spin")} aria-hidden /> : <RefreshCw className="size-3.5" aria-hidden />}
            Re-run answer
          </motion.button>
          <span id={`${uid}-hint`} className="sr-only">
            Regenerates the answer using only the switched-on chunks
          </span>
        </div>
      </div>
      <style>{`@keyframes rag-scan{0%{top:-12px}100%{top:36px}}@keyframes rag-shim{0%{transform:translateX(-100%)}100%{transform:translateX(100%)}}`}</style>
    </div>
  );
}

export default function RagRetrievalDemo() {
  return (
    <div className="flex h-[380px] w-full items-start justify-center pt-2">
      <RagRetrieval autoPlay />
    </div>
  );
}