Trace Flamegraph

"use client";

import * as React from "react";
import { useInView, useReducedMotion } from "motion/react";
import { cn } from "@/lib/utils";

export interface Span {
  name: string;
  service: string;
  /** Start offset from trace root, in ms. */
  start: number;
  /** Duration in ms. */
  duration: number;
  /** Mark this span as on the critical path (longest-latency chain). */
  critical?: boolean;
  /** Marks the span as an error (renders with the destructive tone). */
  error?: boolean;
  children?: Span[];
}

export interface TraceFlamegraphProps
  extends Omit<React.ComponentProps<"div">, "children"> {
  /** Root span of the trace. Children nest by depth; all share one time scale. */
  root?: Span;
  /** Accent color for the shared time axis + critical-path emphasis. */
  accent?: string;
  /** Emphasize the critical path; non-critical spans dim. */
  showCriticalPath?: boolean;
  /** Opt out of the enter reveal. */
  static?: boolean;
}

interface FlatSpan extends Span {
  depth: number;
  key: string;
  parentKey: string | null;
  /** All ancestor keys, root-first. */
  ancestors: string[];
}

function flatten(
  span: Span,
  depth = 0,
  key = "0",
  parentKey: string | null = null,
  ancestors: string[] = [],
): FlatSpan[] {
  const self: FlatSpan = { ...span, depth, key, parentKey, ancestors };
  const kids = (span.children ?? []).flatMap((c, i) =>
    flatten(c, depth + 1, `${key}.${i}`, key, [...ancestors, key]),
  );
  return [self, ...kids];
}

// Deterministic hue per service, drawn from the semantic palette.
const SERVICE_TONES = [
  "var(--color-primary)",
  "var(--color-success)",
  "var(--color-warning)",
  "var(--color-muted-foreground)",
  "var(--color-destructive)",
];

/** Compute "nice" axis tick offsets (ms) across [0, total]. */
function computeTicks(total: number, target = 5): number[] {
  const raw = total / target;
  const mag = Math.pow(10, Math.floor(Math.log10(raw)));
  const norm = raw / mag;
  const nice = norm >= 5 ? 10 : norm >= 2 ? 5 : norm >= 1 ? 2 : 1;
  const step = nice * mag;
  const ticks: number[] = [];
  for (let t = 0; t <= total + 1e-6; t += step) ticks.push(Math.round(t));
  return ticks;
}

const SCENARIOS: Record<string, Span> = {
  fast: {
    name: "POST /v1/checkout",
    service: "api-gateway",
    start: 0,
    duration: 214,
    critical: true,
    children: [
      { name: "auth.verify", service: "auth-svc", start: 3, duration: 22 },
      {
        name: "orders.create",
        service: "orders-svc",
        start: 27,
        duration: 168,
        critical: true,
        children: [
          { name: "SELECT inventory", service: "postgres", start: 31, duration: 34 },
          {
            name: "payments.charge",
            service: "billing-svc",
            start: 68,
            duration: 118,
            critical: true,
            children: [
              { name: "POST stripe.com", service: "external", start: 76, duration: 96, critical: true },
            ],
          },
        ],
      },
      { name: "cache.write", service: "redis", start: 198, duration: 12 },
    ],
  },
  slow: {
    name: "POST /v1/checkout",
    service: "api-gateway",
    start: 0,
    duration: 742,
    critical: true,
    children: [
      { name: "auth.verify", service: "auth-svc", start: 4, duration: 41 },
      {
        name: "orders.create",
        service: "orders-svc",
        start: 48,
        duration: 640,
        critical: true,
        children: [
          { name: "SELECT inventory", service: "postgres", start: 54, duration: 96 },
          {
            name: "payments.charge",
            service: "billing-svc",
            start: 156,
            duration: 470,
            critical: true,
            children: [
              { name: "POST stripe.com", service: "external", start: 172, duration: 428, critical: true },
            ],
          },
        ],
      },
      { name: "cache.write", service: "redis", start: 700, duration: 30 },
    ],
  },
  error: {
    name: "POST /v1/checkout",
    service: "api-gateway",
    start: 0,
    duration: 528,
    critical: true,
    error: true,
    children: [
      { name: "auth.verify", service: "auth-svc", start: 4, duration: 39 },
      {
        name: "orders.create",
        service: "orders-svc",
        start: 47,
        duration: 456,
        critical: true,
        error: true,
        children: [
          { name: "SELECT inventory", service: "postgres", start: 52, duration: 71 },
          {
            name: "payments.charge",
            service: "billing-svc",
            start: 128,
            duration: 372,
            critical: true,
            error: true,
            children: [
              { name: "POST stripe.com", service: "external", start: 144, duration: 348, critical: true, error: true },
            ],
          },
        ],
      },
    ],
  },
};

const DEFAULT_ROOT = SCENARIOS.slow;

export function TraceFlamegraph({
  root = DEFAULT_ROOT,
  accent,
  showCriticalPath = false,
  static: isStatic = false,
  className,
  style,
  ...props
}: TraceFlamegraphProps) {
  const [active, setActive] = React.useState<string | null>(null);
  const reduced = useReducedMotion() ?? false;
  const ref = React.useRef<HTMLDivElement>(null);
  const inView = useInView(ref, { once: true, margin: "0px 0px -32px 0px" });
  const reveal = isStatic || reduced || inView;

  const flat = React.useMemo(() => flatten(root), [root]);
  const total = root.duration || 1;
  const traceStart = root.start;
  const maxDepth = Math.max(...flat.map((s) => s.depth));

  const services = React.useMemo(
    () => Array.from(new Set(flat.map((s) => s.service))),
    [flat],
  );
  const toneFor = (s: FlatSpan) => {
    if (s.error) return "var(--color-destructive)";
    // When the critical path is emphasized, its spans adopt the accent so the
    // color control has a visible effect on the highlighted chain.
    if (showCriticalPath && s.critical) return "var(--pg-flame-accent)";
    return SERVICE_TONES[services.indexOf(s.service) % SERVICE_TONES.length];
  };

  const ticks = React.useMemo(() => computeTicks(total), [total]);

  // Layout in a single coordinate space. x is a percentage of the shared
  // time scale so every span rect aligns; rows align on a fixed row height.
  const rowH = 26;
  const gap = 3;
  const xPct = (t: number) => ((t - traceStart) / total) * 100;

  const activeSpan = active ? flat.find((f) => f.key === active) ?? null : null;
  const relatedKeys = React.useMemo(() => {
    if (!activeSpan) return null;
    const set = new Set<string>([activeSpan.key, ...activeSpan.ancestors]);
    // descendants: any span whose ancestors include the active key
    for (const s of flat) if (s.ancestors.includes(activeSpan.key)) set.add(s.key);
    return set;
  }, [activeSpan, flat]);

  return (
    <div
      ref={ref}
      className={cn(
        "w-full max-w-xl rounded-xl bg-card p-4 text-card-foreground shadow-border",
        className,
      )}
      style={
        {
          "--pg-flame-accent": accent ?? "var(--color-primary)",
          ...style,
        } as React.CSSProperties
      }
      {...props}
    >
      <div className="mb-2.5 flex items-center justify-between gap-3">
        <div className="flex min-w-0 items-center gap-2">
          <span
            className={cn(
              "size-1.5 shrink-0 rounded-full",
              root.error ? "bg-destructive" : "bg-success",
            )}
          />
          <span className="truncate font-mono text-xs font-medium">
            {root.name}
          </span>
        </div>
        <span className="shrink-0 text-xs tabular-nums text-muted-foreground">
          {total.toFixed(0)}ms · {flat.length} spans
        </span>
      </div>

      {/* time axis */}
      <div className="relative mb-1 h-4 select-none">
        {ticks.map((t) => {
          const pct = xPct(t);
          return (
            <span
              key={t}
              className="absolute top-0 -translate-x-1/2 text-[10px] tabular-nums text-muted-foreground first:translate-x-0 last:-translate-x-full"
              style={{ left: `${pct}%` }}
            >
              {t}ms
            </span>
          );
        })}
      </div>

      <div
        className="relative"
        style={{ height: (maxDepth + 1) * (rowH + gap) - gap }}
        onMouseLeave={() => setActive(null)}
      >
        {/* gridlines aligned to the axis ticks */}
        <div className="pointer-events-none absolute inset-0" aria-hidden>
          {ticks.map((t) => (
            <span
              key={t}
              className="absolute top-0 bottom-0 w-px bg-border/60"
              style={{ left: `${xPct(t)}%` }}
            />
          ))}
        </div>

        {flat.map((s, i) => {
          const left = xPct(s.start);
          const width = Math.max((s.duration / total) * 100, 0.6);
          const tone = toneFor(s);
          const isActive = active === s.key;
          const related = relatedKeys?.has(s.key) ?? true;
          const dimForPath = showCriticalPath && !s.critical;
          const dimForHover = active != null && !related;
          const dim = dimForHover || dimForPath;
          const emphasize = showCriticalPath && s.critical;
          const pct = ((s.duration / total) * 100).toFixed(1);
          return (
            <button
              key={s.key}
              type="button"
              aria-label={`${s.name}, ${s.service}, ${s.duration}ms, ${pct}% of trace`}
              aria-pressed={isActive}
              onMouseEnter={() => setActive(s.key)}
              onFocus={() => setActive(s.key)}
              onBlur={() => setActive((a) => (a === s.key ? null : a))}
              className={cn(
                "absolute overflow-hidden rounded-[4px] px-1.5 text-left outline-none",
                "transition-[opacity,filter,translate] duration-150 ease-out",
                "focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1 focus-visible:ring-offset-card",
                !isStatic &&
                  "active:not-disabled:scale-[0.99] transition-[opacity,filter,translate,scale]",
              )}
              style={{
                left: `${left}%`,
                width: `${width}%`,
                top: s.depth * (rowH + gap),
                height: rowH,
                background: `color-mix(in oklch, ${tone} ${
                  emphasize ? 34 : 24
                }%, var(--color-secondary))`,
                boxShadow: `inset 0 0 0 1px color-mix(in oklch, ${tone} ${
                  emphasize ? 80 : 55
                }%, transparent)`,
                opacity: reveal ? (dim ? 0.4 : 1) : 0,
                translate: reveal ? "0 0" : "0 8px",
                filter: reveal ? "blur(0)" : "blur(4px)",
                transitionDelay: reveal && !reduced ? `${Math.min(i * 22, 260)}ms` : "0ms",
              }}
            >
              <span className="flex h-full items-center gap-1 truncate text-[11px] font-medium leading-none">
                {s.error && (
                  <span className="size-1 shrink-0 rounded-full bg-destructive" />
                )}
                <span className="truncate">{s.name}</span>
              </span>
            </button>
          );
        })}
      </div>

      {/* inspector */}
      <div className="mt-2.5 flex min-h-[36px] items-center rounded-lg border border-border bg-background px-2.5 py-1.5 text-[11px]">
        {activeSpan ? (
          <div className="flex w-full flex-wrap items-center gap-x-3 gap-y-0.5">
            <span
              className="size-2 shrink-0 rounded-full"
              style={{ background: toneFor(activeSpan) }}
            />
            <span className="min-w-0 truncate font-mono font-medium">
              {activeSpan.name}
            </span>
            <span className="truncate text-muted-foreground">
              {activeSpan.service}
            </span>
            <span className="ml-auto shrink-0 tabular-nums text-muted-foreground">
              +{activeSpan.start - traceStart}ms
            </span>
            <span className="shrink-0 font-medium tabular-nums">
              {activeSpan.duration}ms
            </span>
            <span className="shrink-0 tabular-nums text-muted-foreground">
              {((activeSpan.duration / total) * 100).toFixed(1)}%
            </span>
          </div>
        ) : (
          <span className="text-muted-foreground">
            Hover a span to inspect timing.
          </span>
        )}
      </div>
    </div>
  );
}