Thought Trace

"use client"

import * as React from "react"
import { ChevronRightIcon, GlobeIcon, LightbulbIcon, SearchIcon, TerminalIcon } from "lucide-react"
import { AnimatePresence, motion } from "motion/react"

import { cn } from "@/lib/utils"
import { useControllableState } from "@/hooks/use-controllable-state"
import { duration, ease, spring, useMotionSafe } from "@/lib/motion"
import { TextShimmer } from "@/components/ui/text-shimmer"

export type TraceStepKind = "thought" | "search" | "read" | "tool"

export type TraceStep = {
  id?: string
  /** What kind of step: a line of reasoning, a web search, a page or file read, or a tool call. */
  kind?: TraceStepKind
  text: string
  /** A second line, such as the query, the URL or the command. */
  detail?: string
}

export type ThoughtTraceProps = Omit<React.ComponentProps<"div">, "children"> & {
  /** "thinking" while the model reasons; "done" once it answers. */
  status: "thinking" | "done"
  steps: TraceStep[]
  /** How long it thought, in seconds. Measured from the first render to "done" when left out. */
  seconds?: number
  open?: boolean
  defaultOpen?: boolean
  onOpenChange?: (open: boolean) => void
  /** Replaces the lightbulb. */
  icon?: React.ReactNode
  /** Header text while thinking. */
  thinkingLabel?: string
  /** Header text when done. */
  doneLabel?: (seconds: number) => string
}

const kindIcons = {
  thought: null,
  search: SearchIcon,
  read: GlobeIcon,
  tool: TerminalIcon,
} as const

/**
 * The reasoning behind a reply, folded into one line. It shimmers “Thinking” while the model
 * works, then says how long it took; opening it lists each step, including searches, pages read
 * and tool calls.
 */
function ThoughtTrace({
  status,
  steps,
  seconds: secondsProp,
  open: openProp,
  defaultOpen = false,
  onOpenChange,
  icon,
  thinkingLabel = "Thinking",
  doneLabel = (s) => `Thought for ${s < 1 ? "a moment" : `${s}s`}`,
  className,
  ...props
}: ThoughtTraceProps) {
  const [open, setOpen] = useControllableState({ value: openProp, defaultValue: defaultOpen, onChange: onOpenChange })
  const thinking = status === "thinking"
  const started = React.useRef<number | null>(null)
  const [measured, setMeasured] = React.useState<number>()
  const listId = React.useId()
  const { reduced, safe } = useMotionSafe()

  // Time the thinking when the caller does not.
  React.useEffect(() => {
    const now = performance.now()
    if (thinking) started.current ??= now
    const begun = started.current
    if (!thinking) started.current = null
    const frame = requestAnimationFrame(() =>
      setMeasured(thinking || begun === null ? undefined : Math.round((now - begun) / 1000))
    )
    return () => cancelAnimationFrame(frame)
  }, [thinking])

  const seconds = secondsProp ?? measured ?? 0

  return (
    <div data-slot="thought-trace" data-status={status} className={cn("flex w-full flex-col", className)} {...props}>
      <button
        type="button"
        aria-expanded={open}
        aria-controls={listId}
        onClick={() => setOpen(!open)}
        className="group -ml-1 flex w-fit cursor-pointer items-center gap-2 rounded-md px-1 py-1 text-sm text-muted-foreground outline-none transition-colors hover:text-foreground focus-visible:ring-2 focus-visible:ring-ring"
      >
        <span aria-hidden="true" className="grid size-4 place-items-center [&_svg]:size-4">
          {icon ?? <LightbulbIcon className={cn(thinking && "text-foreground")} />}
        </span>
        <span aria-live="polite" className="font-medium">
          {thinking ? (
            <TextShimmer duration={1.8}>{thinkingLabel}</TextShimmer>
          ) : (
            doneLabel(seconds)
          )}
        </span>
        <motion.span
          aria-hidden="true"
          animate={{ rotate: open ? 90 : 0 }}
          transition={safe(spring.snappy)}
          className="grid place-items-center"
        >
          <ChevronRightIcon className="size-3.5" />
        </motion.span>
      </button>

      <AnimatePresence initial={false}>
        {open && (
          <motion.div
            id={listId}
            initial={{ height: 0, opacity: 0 }}
            animate={{ height: "auto", opacity: 1 }}
            exit={{ height: 0, opacity: 0 }}
            transition={safe({ duration: duration.slow, ease: ease.out })}
            className="overflow-hidden"
          >
            <ol aria-label="Reasoning steps" className="relative mt-1 flex flex-col gap-3 py-1.5">
              <AnimatePresence initial={false}>
                {steps.map((step, index) => {
                  const Icon = kindIcons[step.kind ?? "thought"]
                  return (
                    <motion.li
                      key={step.id ?? index}
                      initial={{ opacity: 0, y: reduced ? 0 : 6 }}
                      animate={{ opacity: 1, y: 0 }}
                      transition={safe({ duration: duration.base, ease: ease.out })}
                      className="relative grid grid-cols-[1rem_minmax(0,1fr)] gap-x-3 gap-y-0.5 text-[13.5px] leading-6 before:absolute before:top-3 before:bottom-[-0.75rem] before:left-[7.5px] before:w-px before:bg-border last:before:hidden"
                    >
                      <span
                        aria-hidden="true"
                        className="relative z-10 col-start-1 row-span-2 mt-1 grid size-4 place-items-center rounded-full bg-background text-muted-foreground [&_svg]:size-3"
                      >
                        {Icon ? <Icon /> : <span className="size-1.5 rounded-full bg-muted-foreground/60" />}
                      </span>
                      <span className={cn("col-start-2", step.kind && step.kind !== "thought" ? "text-foreground" : "text-muted-foreground")}>
                        {step.text}
                      </span>
                      {step.detail && (
                        <span className="col-start-2 truncate font-mono text-[12px] text-muted-foreground">{step.detail}</span>
                      )}
                    </motion.li>
                  )
                })}
              </AnimatePresence>
              {thinking && (
                <li className="pl-7 text-[13.5px] leading-6 text-muted-foreground">
                  <span className="sr-only">Still thinking</span>
                  <span aria-hidden="true" className="inline-flex gap-1">
                    {[0, 1, 2].map((dot) => (
                      <span
                        key={dot}
                        className="size-1 rounded-full bg-current motion-safe:animate-pulse"
                        style={{ animationDelay: `${dot * 160}ms` }}
                      />
                    ))}
                  </span>
                </li>
              )}
            </ol>
          </motion.div>
        )}
      </AnimatePresence>
    </div>
  )
}

export { ThoughtTrace }