Trace waterfallNew

"use client";

import { useEffect, useId, useRef, useState } from "react";
import { Activity, AlertTriangle, Bot, Database, Sparkles, Wrench } from "lucide-react";
import { AnimatePresence, motion, useReducedMotion } from "motion/react";
import { cn } from "@/lib/utils";

export type TraceSpanKind = "agent" | "llm" | "tool" | "retrieval";

export type TraceSpan = {
  id: string;
  name: string;
  kind: TraceSpanKind;
  /** Start offset from the beginning of the trace, in ms. */
  start: number;
  /** Duration in ms. */
  duration: number;
  depth?: number;
  error?: string;
  tokens?: number;
  /** Short input/output preview shown in the detail panel. */
  input?: string;
  output?: string;
};

export type TraceWaterfallProps = {
  traceId?: string;
  spans?: TraceSpan[];
  /** Elapsed time of a live trace in ms. Leave unset for a finished trace. */
  now?: number;
  defaultSelected?: string;
  onSelect?: (span: TraceSpan) => void;
  /** Loop a scripted live run, then inspect the failed span. */
  autoPlay?: boolean;
  className?: string;
};

const DEFAULT_SPANS: TraceSpan[] = [
  { id: "s1", name: "agent.run", kind: "agent", start: 0, duration: 3100, depth: 0, tokens: 4870, input: "“Why was Acme billed twice in October?”", output: "Answer streamed · 2 citations" },
  { id: "s2", name: "plan", kind: "llm", start: 40, duration: 560, depth: 1, tokens: 1240, input: "system + 3 tools · 1.1k tok", output: "3 steps: search → fetch invoice → summarize" },
  { id: "s3", name: "search_docs", kind: "retrieval", start: 640, duration: 470, depth: 1, input: "q=“duplicate charge october”", output: "6 chunks · top score 0.82" },
  { id: "s4", name: "billing-agent", kind: "agent", start: 1160, duration: 760, depth: 1, tokens: 410, input: "handoff · invoice lookup", output: "INV-2084 · 2 charges found" },
  { id: "s5", name: "fetch_invoice", kind: "tool", start: 1200, duration: 300, depth: 2, error: "504 Gateway Timeout", input: "GET /invoices/INV-2084", output: "504 after 300ms · retrying" },
  { id: "s6", name: "fetch_invoice", kind: "tool", start: 1540, duration: 340, depth: 2, input: "GET /invoices/INV-2084 (retry 1)", output: "200 · 2 line items" },
  { id: "s7", name: "summarize", kind: "llm", start: 1960, duration: 1110, depth: 1, tokens: 3220, input: "6 chunks + invoice · 2.9k tok", output: "312 tok · 288 tok/s" },
];

const KIND = {
  agent: { icon: Bot, bar: "bg-emerald-500", text: "text-emerald-600 dark:text-emerald-400", label: "Agent" },
  llm: { icon: Sparkles, bar: "bg-violet-500", text: "text-violet-600 dark:text-violet-400", label: "LLM" },
  tool: { icon: Wrench, bar: "bg-amber-500", text: "text-amber-600 dark:text-amber-400", label: "Tool" },
  retrieval: { icon: Database, bar: "bg-sky-500", text: "text-sky-600 dark:text-sky-400", label: "Retrieval" },
} as const;

const EASE = [0.22, 1, 0.36, 1] as const;
const TICK = 100;
const LOOP = 9000;
const START = 300;

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 };
}

const fmt = (ms: number) => (ms >= 1000 ? `${(ms / 1000).toFixed(ms >= 10000 ? 0 : 1)}s` : `${Math.round(ms)}ms`);

/** Live span waterfall for an agent run: LLM, tool, retrieval and sub-agent spans on one timeline, with a detail panel per span. */
export function TraceWaterfall({ traceId = "run_7f3a", spans = DEFAULT_SPANS, now: nowProp, defaultSelected, onSelect, autoPlay = false, className }: TraceWaterfallProps) {
  const reduce = useReducedMotion() ?? false;
  const uid = useId();
  const [touched, setTouched] = useState(false);
  const [picked, setPicked] = useState<string | null>(defaultSelected ?? null);
  const demo = autoPlay && !reduce && !touched && nowProp === undefined;
  const { ref, t } = useDemoClock(demo, LOOP);

  const total = Math.max(...spans.map((s) => s.start + s.duration));
  const now = demo ? Math.max(0, Math.min(total, t - START)) : (nowProp ?? total);
  const live = now < total;
  const selected = demo ? (t >= 4300 ? "s5" : null) : picked;
  const sel = spans.find((s) => s.id === selected);
  const errors = spans.filter((s) => s.error && now >= s.start + s.duration).length;
  const tokens = spans.filter((s) => (s.depth ?? 0) > 0 && now >= s.start + s.duration).reduce((a, s) => a + (s.tokens ?? 0), 0);

  const scaleMax = Math.ceil(total / 1000) * 1000;
  const pct = (ms: number) => (ms / scaleMax) * 100;
  const ticks = Array.from({ length: scaleMax / 1000 + 1 }, (_, i) => i * 1000);

  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 rounded-xl border bg-background shadow-sm">
            <Activity className="size-4" aria-hidden />
            {live && !reduce && <span className="absolute -top-0.5 -right-0.5 size-2 animate-ping rounded-full bg-sky-500" aria-hidden />}
            {live && <span className="absolute -top-0.5 -right-0.5 size-2 rounded-full bg-sky-500" aria-hidden />}
          </span>
          <div className="min-w-0 flex-1">
            <p className="text-sm font-semibold">Trace</p>
            <p className="truncate font-mono text-[10.5px] text-muted-foreground">
              {traceId} · {spans.filter((s) => now >= s.start).length} spans · {tokens.toLocaleString()} tok
              {errors > 0 && <span className="text-rose-600 dark:text-rose-400"> · {errors} error</span>}
            </p>
          </div>
          <span
            className={cn(
              "inline-flex shrink-0 items-center gap-1.5 rounded-full border px-2 py-0.5 font-mono text-[10px] tabular-nums",
              live ? "border-sky-500/30 bg-sky-500/10 text-sky-700 dark:text-sky-300" : "bg-background text-muted-foreground",
            )}
            aria-live="polite"
          >
            {live && <span className={cn("size-1.5 rounded-full bg-sky-500", !reduce && "animate-pulse")} />}
            {live ? `Live ${fmt(now)}` : `Done ${fmt(total)}`}
          </span>
        </div>

        {/* Ruler */}
        <div className="flex items-end px-3 pb-1">
          <span className="w-[132px] shrink-0 pl-1 text-[10px] text-muted-foreground">Span</span>
          <div className="relative h-4 flex-1">
            {ticks.map((ms) => (
              <span key={ms} className="absolute bottom-0 -translate-x-1/2 font-mono text-[9px] text-muted-foreground tabular-nums first:translate-x-0 last:-translate-x-full" style={{ left: `${pct(ms)}%` }}>
                {ms / 1000}s
              </span>
            ))}
          </div>
        </div>

        {/* Rows */}
        <div className="relative px-3 pb-2" aria-label="Spans">
          <div className="pointer-events-none absolute inset-y-0 right-3 left-[144px]" aria-hidden>
            {ticks.map((ms) => (
              <span key={ms} className="absolute inset-y-0 w-px bg-border/60" style={{ left: `${pct(ms)}%` }} />
            ))}
            {live && (
              <motion.span
                className="absolute inset-y-0 w-px bg-sky-500"
                initial={false}
                animate={{ left: `${pct(now)}%` }}
                transition={{ duration: TICK / 1000, ease: "linear" }}
              >
                <span className="absolute -top-0.5 -left-[3px] size-[7px] rounded-full bg-sky-500 shadow-[0_0_0_3px_rgb(14_165_233/0.2)]" />
              </motion.span>
            )}
          </div>
          {spans.map((s) => {
            const started = now >= s.start;
            const end = Math.min(now, s.start + s.duration);
            const running = started && now < s.start + s.duration;
            const failed = !!s.error && !running && started;
            const k = KIND[s.kind];
            const Icon = failed ? AlertTriangle : k.icon;
            const isSel = s.id === selected;
            return (
              <button
                key={s.id}
                type="button"
                aria-pressed={isSel}
                aria-label={`${s.name}, ${k.label}, ${fmt(s.duration)}${s.error ? `, failed: ${s.error}` : ""}`}
                onClick={() => {
                  setTouched(true);
                  setPicked(isSel ? null : s.id);
                  if (!isSel) onSelect?.(s);
                }}
                className={cn(
                  "relative flex h-7 w-full cursor-pointer items-center rounded-md text-left outline-none transition-[background-color,opacity] duration-300 focus-visible:ring-2 focus-visible:ring-ring",
                  isSel ? "bg-muted" : "hover:bg-muted/50",
                  !started && "opacity-35",
                )}
              >
                <span className="flex w-[132px] shrink-0 items-center gap-1.5 pr-2" style={{ paddingLeft: 4 + (s.depth ?? 0) * 10 }}>
                  {(s.depth ?? 0) > 0 && <span className="h-px w-1.5 shrink-0 bg-border" aria-hidden />}
                  <Icon className={cn("size-3 shrink-0", failed ? "text-rose-500" : k.text)} aria-hidden />
                  <span className={cn("truncate font-mono text-[10.5px]", (s.depth ?? 0) === 0 && "font-semibold")}>{s.name}</span>
                </span>
                <span className="relative h-full flex-1">
                  {started && (
                    <motion.span
                      className={cn(
                        "absolute top-1/2 h-3.5 -translate-y-1/2 overflow-hidden rounded-[4px]",
                        failed ? "bg-rose-500" : k.bar,
                        (s.depth ?? 0) === 0 && !failed && "opacity-35",
                        isSel && "ring-2 ring-foreground/70 ring-offset-1 ring-offset-card",
                      )}
                      style={{ left: `${pct(s.start)}%` }}
                      initial={false}
                      animate={{ width: `${Math.max(pct(end - s.start), 0.8)}%` }}
                      transition={{ duration: demo ? TICK / 1000 : 0.5, ease: demo ? "linear" : EASE }}
                    >
                      {running && !reduce && <span className="absolute inset-0 animate-[tw-stripe_0.8s_linear_infinite] bg-[repeating-linear-gradient(-45deg,rgb(255_255_255/0.35)_0_4px,transparent_4px_8px)] bg-[length:11.3px_100%]" />}
                    </motion.span>
                  )}
                  {started && !running && (s.depth ?? 0) > 0 && (
                    <motion.span
                      initial={reduce ? false : { opacity: 0 }}
                      animate={{ opacity: 1 }}
                      className={cn("absolute top-1/2 -translate-y-1/2 pl-1 font-mono text-[9px] whitespace-nowrap tabular-nums", failed ? "text-rose-600 dark:text-rose-400" : "text-muted-foreground")}
                      style={{ left: `${pct(s.start + s.duration)}%`, ...(pct(s.start + s.duration) > 80 ? { left: "auto", right: `${100 - pct(s.start)}%`, paddingRight: 4 } : {}) }}
                    >
                      {fmt(s.duration)}
                    </motion.span>
                  )}
                </span>
              </button>
            );
          })}
        </div>

        {/* Detail */}
        <AnimatePresence initial={false} mode="popLayout">
          {sel ? (
            <motion.div
              key={sel.id}
              id={`${uid}-detail`}
              initial={reduce ? false : { opacity: 0, y: 8 }}
              animate={{ opacity: 1, y: 0 }}
              exit={reduce ? undefined : { opacity: 0, y: -6 }}
              transition={{ duration: 0.3, ease: EASE }}
              className="border-t px-4 py-3"
            >
              <div className="mb-2 flex items-center gap-2">
                <span className={cn("rounded-full px-1.5 py-px text-[10px] font-medium", sel.error ? "bg-rose-500/10 text-rose-700 dark:text-rose-300" : "bg-muted text-muted-foreground")}>
                  {sel.error ? "Error" : KIND[sel.kind].label}
                </span>
                <span className="font-mono text-xs font-medium">{sel.name}</span>
                <span className="flex-1" />
                <span className="font-mono text-[10px] text-muted-foreground tabular-nums">
                  {fmt(sel.duration)}
                  {sel.tokens ? ` · ${sel.tokens.toLocaleString()} tok` : ""}
                </span>
              </div>
              <dl className="grid grid-cols-[auto_1fr] gap-x-3 gap-y-1 rounded-lg bg-muted/60 px-2.5 py-2 font-mono text-[10.5px]">
                <dt className="text-muted-foreground">in</dt>
                <dd className="truncate">{sel.input ?? "—"}</dd>
                <dt className="text-muted-foreground">out</dt>
                <dd className={cn("truncate", sel.error && "text-rose-600 dark:text-rose-400")}>{sel.error ? `${sel.error} · ${sel.output ?? ""}` : (sel.output ?? "—")}</dd>
              </dl>
            </motion.div>
          ) : (
            <motion.div
              key="legend"
              initial={reduce ? false : { opacity: 0 }}
              animate={{ opacity: 1 }}
              exit={reduce ? undefined : { opacity: 0 }}
              className="flex flex-wrap items-center gap-x-3 gap-y-1 border-t px-4 py-2.5 text-[10.5px] text-muted-foreground"
            >
              {(Object.keys(KIND) as TraceSpanKind[]).map((key) => (
                <span key={key} className="inline-flex items-center gap-1.5">
                  <span className={cn("size-2 rounded-[3px]", KIND[key].bar)} />
                  {KIND[key].label}
                </span>
              ))}
              <span className="ml-auto">Click a span to inspect</span>
            </motion.div>
          )}
        </AnimatePresence>
      </div>
      <style>{`@keyframes tw-stripe{to{background-position:11.3px 0}}`}</style>
    </div>
  );
}

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