Trace Viewer

"use client";

import * as React from "react";
import { CheckCircleIcon, ClockIcon, CpuChipIcon, WrenchScrewdriverIcon, XCircleIcon } from "@heroicons-animated/react";

import { cn } from "@/lib/utils";
import { Badge } from "@/components/ui/badge";
import { CodeBlock, CodeBlockContent, CodeBlockCopyButton, CodeBlockGroup, CodeBlockHeader, CodeBlockShiki } from "@/components/pandacoderz-ui/code-block";

type SpanKind = "llm" | "tool" | "retrieval" | "guard";
type Span = { id: string; name: string; kind: SpanKind; start: number; duration: number; depth: number; status: "ok" | "error"; tokensIn?: number; tokensOut?: number; input: unknown; output: unknown };
type Trace = { id: string; model: string; status: "ok" | "error"; latency: number; tokensIn: number; tokensOut: number; cost: number; at: string; spans: Span[] };

const traces: Trace[] = [
  {
    id: "tr_8f2a", model: "claude-sonnet-5", status: "ok", latency: 3120, tokensIn: 1840, tokensOut: 412, cost: 0.0117, at: "14:02:11",
    spans: [
      { id: "s1", name: "chat.completion", kind: "llm", start: 0, duration: 3120, depth: 0, status: "ok", tokensIn: 1840, tokensOut: 412, input: { messages: 6, model: "claude-sonnet-5" }, output: { stop_reason: "end_turn" } },
      { id: "s2", name: "guard.input", kind: "guard", start: 20, duration: 90, depth: 1, status: "ok", input: { policy: "default" }, output: { flagged: false } },
      { id: "s3", name: "retrieval.search", kind: "retrieval", start: 130, duration: 410, depth: 1, status: "ok", input: { query: "astro islands hydration", k: 5 }, output: { hits: 5, topScore: 0.87 } },
      { id: "s4", name: "llm.generate", kind: "llm", start: 560, duration: 1380, depth: 1, status: "ok", tokensIn: 1840, tokensOut: 96, input: { temperature: 0.7 }, output: { tool_use: "search_docs" } },
      { id: "s5", name: "tool.search_docs", kind: "tool", start: 1960, duration: 320, depth: 1, status: "ok", input: { query: "client directives", limit: 2 }, output: { results: 2 } },
      { id: "s6", name: "llm.generate", kind: "llm", start: 2300, duration: 790, depth: 1, status: "ok", tokensIn: 2010, tokensOut: 316, input: { temperature: 0.7 }, output: { stop_reason: "end_turn" } },
    ],
  },
  {
    id: "tr_7c19", model: "claude-haiku-4-5", status: "ok", latency: 880, tokensIn: 620, tokensOut: 140, cost: 0.0011, at: "14:01:47",
    spans: [
      { id: "s1", name: "chat.completion", kind: "llm", start: 0, duration: 880, depth: 0, status: "ok", tokensIn: 620, tokensOut: 140, input: { messages: 2 }, output: { stop_reason: "end_turn" } },
      { id: "s2", name: "guard.input", kind: "guard", start: 10, duration: 60, depth: 1, status: "ok", input: {}, output: { flagged: false } },
      { id: "s3", name: "llm.generate", kind: "llm", start: 80, duration: 780, depth: 1, status: "ok", tokensIn: 620, tokensOut: 140, input: { temperature: 0.3 }, output: { stop_reason: "end_turn" } },
    ],
  },
  {
    id: "tr_6b02", model: "claude-opus-5", status: "error", latency: 5210, tokensIn: 3200, tokensOut: 0, cost: 0.048, at: "13:58:03",
    spans: [
      { id: "s1", name: "chat.completion", kind: "llm", start: 0, duration: 5210, depth: 0, status: "error", tokensIn: 3200, input: { messages: 14 }, output: { error: "tool_timeout" } },
      { id: "s2", name: "retrieval.search", kind: "retrieval", start: 40, duration: 520, depth: 1, status: "ok", input: { query: "pricing tiers" }, output: { hits: 3 } },
      { id: "s3", name: "llm.generate", kind: "llm", start: 580, duration: 1600, depth: 1, status: "ok", tokensIn: 3200, tokensOut: 40, input: {}, output: { tool_use: "fetch_pricing" } },
      { id: "s4", name: "tool.fetch_pricing", kind: "tool", start: 2200, duration: 3000, depth: 1, status: "error", input: { url: "/api/pricing" }, output: { error: "timeout after 3000ms" } },
    ],
  },
];

const kindColor: Record<SpanKind, string> = {
  llm: "bg-brand",
  tool: "bg-chart-2",
  retrieval: "bg-chart-3",
  guard: "bg-chart-4",
};
const kindIcon: Record<SpanKind, React.ComponentType<{ size?: number; className?: string }>> = {
  llm: CpuChipIcon,
  tool: WrenchScrewdriverIcon,
  retrieval: ClockIcon,
  guard: CheckCircleIcon,
};

export type TraceViewerProps = { className?: string };

export default function TraceViewer({ className }: TraceViewerProps) {
  const [traceId, setTraceId] = React.useState(traces[0].id);
  const trace = traces.find((t) => t.id === traceId) ?? traces[0];
  const [spanId, setSpanId] = React.useState(trace.spans[0].id);
  const span = trace.spans.find((s) => s.id === spanId) ?? trace.spans[0];
  const total = trace.spans[0].duration;

  return (
    <div data-slot="trace-viewer" className={cn("flex h-full min-h-0 w-full flex-col overflow-hidden rounded-3xl border bg-background shadow-xs", className)}>
      <header className="flex h-12 shrink-0 items-center justify-between gap-3 border-b px-4">
        <div className="flex items-center gap-2 text-sm font-medium">
          <span className="flex size-6 items-center justify-center rounded-md bg-brand text-primary-foreground"><ClockIcon size={14} className="flex" /></span>
          Traces
        </div>
        <div className="hidden items-center gap-3 text-[11px] text-muted-foreground sm:flex">
          {(Object.keys(kindColor) as SpanKind[]).map((k) => <span key={k} className="flex items-center gap-1.5"><span className={cn("size-2 rounded-sm", kindColor[k])} />{k}</span>)}
        </div>
      </header>

      <div className="grid min-h-0 flex-1 lg:grid-cols-[minmax(0,15rem)_minmax(0,1fr)_minmax(0,20rem)]">
        <aside className="min-h-0 overflow-y-auto border-b bg-surface lg:border-r lg:border-b-0">
          <ul className="p-2">
            {traces.map((t) => (
              <li key={t.id}>
                <button type="button" onClick={() => { setTraceId(t.id); setSpanId(t.spans[0].id); }} className={cn("flex w-full flex-col gap-1 rounded-lg px-2.5 py-2 text-left transition-colors", t.id === traceId ? "bg-brand-soft/60 dark:bg-brand-soft/50" : "hover:bg-accent")}>
                  <div className="flex items-center justify-between gap-2 text-xs">
                    <span className="font-mono">{t.id}</span>
                    {t.status === "ok" ? <CheckCircleIcon size={12} className="flex text-emerald-600 dark:text-emerald-400" /> : <XCircleIcon size={12} className="flex text-red-600 dark:text-red-400" />}
                  </div>
                  <div className="flex items-center justify-between gap-2 text-[11px] text-muted-foreground">
                    <span className="truncate">{t.model}</span>
                    <span className="font-mono tabular-nums">{(t.latency / 1000).toFixed(2)}s</span>
                  </div>
                </button>
              </li>
            ))}
          </ul>
        </aside>

        <div className="flex min-h-0 flex-col">
          <div className="grid shrink-0 grid-cols-2 gap-px border-b bg-border sm:grid-cols-4">
            {[
              ["Latency", `${(trace.latency / 1000).toFixed(2)}s`],
              ["Tokens", `${trace.tokensIn.toLocaleString()} → ${trace.tokensOut.toLocaleString()}`],
              ["Cost", `$${trace.cost.toFixed(4)}`],
              ["Tool calls", String(trace.spans.filter((s) => s.kind === "tool").length)],
            ].map(([label, value]) => (
              <div key={label} className="flex flex-col gap-0.5 bg-background px-4 py-2.5">
                <span className="text-[11px] text-muted-foreground">{label}</span>
                <span className="font-mono text-sm tabular-nums">{value}</span>
              </div>
            ))}
          </div>
          <div className="min-h-0 flex-1 overflow-auto p-4">
            <div className="relative min-w-[32rem]">
              <div className="mb-2 flex justify-between font-mono text-[10px] text-muted-foreground">
                {[0, 0.25, 0.5, 0.75, 1].map((f) => <span key={f}>{Math.round(total * f)}ms</span>)}
              </div>
              <ul className="flex flex-col gap-1">
                {trace.spans.map((s) => {
                  const Icon = kindIcon[s.kind];
                  return (
                    <li key={s.id}>
                      <button type="button" onClick={() => setSpanId(s.id)} className={cn("group/span flex h-8 w-full items-center gap-3 rounded-lg px-2 text-left text-xs transition-colors hover:bg-accent", s.id === spanId && "bg-accent")}>
                        <span className="flex w-44 shrink-0 items-center gap-1.5 truncate" style={{ paddingLeft: s.depth * 14 }}>
                          <Icon size={12} className="flex shrink-0 text-muted-foreground" />
                          <span className={cn("truncate font-mono", s.status === "error" && "text-red-600 dark:text-red-400")}>{s.name}</span>
                        </span>
                        <span className="relative h-4 flex-1 rounded bg-muted/40">
                          <span className={cn("absolute inset-y-0 rounded", kindColor[s.kind], s.status === "error" && "bg-red-500", s.id === spanId ? "opacity-100" : "opacity-80")} style={{ left: `${(s.start / total) * 100}%`, width: `${Math.max(0.6, (s.duration / total) * 100)}%` }} />
                        </span>
                        <span className="w-14 shrink-0 text-right font-mono tabular-nums text-muted-foreground">{s.duration}ms</span>
                      </button>
                    </li>
                  );
                })}
              </ul>
            </div>
          </div>
        </div>

        <aside className="flex min-h-0 flex-col border-t bg-surface lg:border-t-0 lg:border-l">
          <div className="flex shrink-0 flex-col gap-2 border-b p-4">
            <div className="flex items-center justify-between gap-2">
              <span className="truncate font-mono text-sm">{span.name}</span>
              <Badge className={cn("border-transparent", span.status === "ok" ? "bg-emerald-500/15 text-emerald-600 dark:text-emerald-400" : "bg-red-500/15 text-red-600 dark:text-red-400")}>{span.status}</Badge>
            </div>
            <dl className="grid grid-cols-2 gap-x-4 gap-y-1 text-xs">
              <dt className="text-muted-foreground">Kind</dt><dd className="text-right">{span.kind}</dd>
              <dt className="text-muted-foreground">Start</dt><dd className="text-right font-mono tabular-nums">{span.start}ms</dd>
              <dt className="text-muted-foreground">Duration</dt><dd className="text-right font-mono tabular-nums">{span.duration}ms</dd>
              {span.tokensIn !== undefined ? <><dt className="text-muted-foreground">Tokens</dt><dd className="text-right font-mono tabular-nums">{span.tokensIn} → {span.tokensOut ?? 0}</dd></> : null}
            </dl>
          </div>
          <div className="flex min-h-0 flex-1 flex-col gap-3 overflow-y-auto p-4">
            {(["input", "output"] as const).map((k) => (
              <CodeBlock key={k}>
                <CodeBlockHeader className="h-8"><CodeBlockGroup><span className="text-[11px] font-medium uppercase tracking-wide">{k}</span></CodeBlockGroup><CodeBlockCopyButton /></CodeBlockHeader>
                <CodeBlockContent className="max-h-48 pb-3 text-xs"><CodeBlockShiki code={JSON.stringify(span[k], null, 2)} language="json" /></CodeBlockContent>
              </CodeBlock>
            ))}
          </div>
        </aside>
      </div>
    </div>
  );
}