app/frontend/src/components/Runs/MetricChart.tsx

import { useQuery } from "@tanstack/react-query"

import type { HistoryPoint } from "@/client"
import { MAX_SERIES, UplotChart } from "@/components/Common/UplotChart"
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select"
import { cn } from "@/lib/utils"
import {
  compareQueryOptions,
  NO_CURVE,
  runMetricsQueryOptions,
  STEP_AXIS,
  shortenRunLabel,
  TIME_AXIS,
} from "./queries"

/** The metrics one run recorded, in the order they are worth offering. */
export function useMetricNames(runId: string | undefined) {
  const { data } = useQuery({
    ...runMetricsQueryOptions(runId ?? "", ""),
    enabled: Boolean(runId),
  })
  const names = new Set<string>()
  for (const point of data ?? []) if (point.name) names.add(point.name)
  return [...names]
}

/**
 * One metric across one or more runs.
 *
 * The endpoint answers in the chart widget's own series shape, so comparing
 * three curves and showing one are the same call with a different id list.
 *
 * `x` is what the readings are plotted against: the step they were recorded
 * at, the seconds since the run's own first one, or another metric of the same
 * runs. Never a clock — even "time" is an elapsed count, which is what makes
 * two runs started hours apart comparable.
 */
export function RunMetricChart({
  ids,
  metric,
  x = STEP_AXIS,
  refreshMs,
  className,
}: {
  ids: string[]
  metric: string
  x?: string
  refreshMs?: number
  className?: string
}) {
  const { data, isPending } = useQuery(
    compareQueryOptions(ids, metric, x, refreshMs),
  )
  const lines = (data?.lines ?? []).slice(0, MAX_SERIES)
  const labels = lines.map((line) => shortenRunLabel(line.label))
  const plots: HistoryPoint[][] = lines.map((line) =>
    (line.points ?? []).map(([at, value]) => ({ ts: at, value })),
  )
  const drawn = plots.reduce((total, plot) => total + plot.length, 0)

  return (
    <div className={cn("flex min-h-0 flex-col gap-2", className ?? "h-64")}>
      <UplotChart
        labels={labels}
        plots={plots}
        xTime={false}
        yLabel={metric}
        pending={isPending}
        empty={NO_CURVE}
      />
      {drawn > 0 && (data?.lines?.length ?? 0) > MAX_SERIES && (
        <p className="text-muted-foreground text-xs">
          Showing {MAX_SERIES} of {data?.lines?.length} runs.
        </p>
      )}
    </div>
  )
}

/**
 * The axis a picker sets, written beside it rather than into every option.
 *
 * The dropdown then lists names — "step", "loss" — instead of repeating "vs"
 * once per row, and which axis is being set is said once, where it belongs.
 */
function AxisLabel({ children }: { children: string }) {
  return (
    <span className="font-medium text-muted-foreground text-xs" aria-hidden>
      {children}
    </span>
  )
}

/** The metric picker both the detail and the comparison sit under. */
export function MetricPicker({
  names,
  value,
  axis,
  onChange,
}: {
  names: string[]
  value: string
  /** Named where the chart has an axis for it to be; bare on its own. */
  axis?: string
  onChange: (name: string) => void
}) {
  if (names.length === 0) return null
  return (
    <div className="flex items-center gap-1.5">
      {axis && <AxisLabel>{axis}</AxisLabel>}
      <Select value={value} onValueChange={onChange}>
        <SelectTrigger
          className={cn("h-8", axis ? "w-52" : "w-56")}
          aria-label={axis ? `${axis} axis` : "Metric"}
        >
          <SelectValue placeholder="Metric" />
        </SelectTrigger>
        <SelectContent>
          {names.map((name) => (
            <SelectItem key={name} value={name}>
              {name}
            </SelectItem>
          ))}
        </SelectContent>
      </Select>
    </div>
  )
}

/** What the comparison is plotted against. */
export function AxisPicker({
  names,
  metric,
  value,
  axis = "X",
  onChange,
}: {
  names: string[]
  /** Excluded from the choices: a metric against itself is a straight line. */
  metric: string
  value: string
  axis?: string
  onChange: (x: string) => void
}) {
  const options = [
    { value: STEP_AXIS, label: "step" },
    { value: TIME_AXIS, label: "time (s)" },
    ...names
      .filter((name) => name !== metric)
      .map((name) => ({ value: name, label: name })),
  ]
  return (
    <div className="flex items-center gap-1.5">
      <AxisLabel>{axis}</AxisLabel>
      <Select value={value} onValueChange={onChange}>
        <SelectTrigger className="h-8 w-44" aria-label={`${axis} axis`}>
          <SelectValue />
        </SelectTrigger>
        <SelectContent>
          {options.map((option) => (
            <SelectItem key={option.value} value={option.value}>
              {option.label}
            </SelectItem>
          ))}
        </SelectContent>
      </Select>
    </div>
  )
}