app/frontend/src/components/Runs/queries.ts

import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
import { useNavigate } from "@tanstack/react-router"
import { useMemo } from "react"

import { OpenAPI, RunsService } from "@/client"
import { flowQueryOptions } from "@/components/Flow/queries"
import { apiToken } from "@/lib/portal"

/**
 * How deep one page of the run list is read.
 *
 * The endpoint caps a page at 500; asking for a hundred at a time keeps the
 * first paint quick and leaves "Load more" something to do.
 */
export const PAGE = 100

/** How far the list will page before it stops offering to go deeper. */
export const LIST_CAP = 500

export type RunFilters = {
  flow?: string
  status?: string
  group?: string
}

export const runKeys = {
  all: ["runs"] as const,
  overview: ["runs", "overview"] as const,
  list: (filters: RunFilters) => ["runs", "list", filters] as const,
  detail: (id: string) => ["runs", "detail", id] as const,
  metrics: (id: string, name: string) =>
    ["runs", "detail", id, "metrics", name] as const,
  compares: ["runs", "compare"] as const,
  compare: (ids: string[], metric: string, x: string) =>
    ["runs", "compare", ids.join(","), metric, x] as const,
}

/** The resting surface these screens are built from, as Health names it. */
export const CARD = "rounded-lg border border-border bg-card p-4 shadow-e1"

/** The statuses a run passes through, as the filter offers them. */
export const STATUSES = [
  "queued",
  "running",
  "ok",
  "error",
  "cancelled",
  "abandoned",
] as const

/** A run that has not settled is still worth re-reading. */
export const isLive = (status: string) =>
  status === "queued" || status === "running"

export const runsInfiniteQueryOptions = (filters: RunFilters) => ({
  queryKey: runKeys.list(filters),
  queryFn: ({ pageParam }: { pageParam: number }) =>
    RunsService.readRuns({ ...filters, limit: PAGE, offset: pageParam }),
  initialPageParam: 0,
  getNextPageParam: (last: unknown[], all: unknown[][]) => {
    const read = all.reduce((total, page) => total + page.length, 0)
    // A short page is the end of the history; the cap is the end of what this
    // list will show of it.
    return last.length < PAGE || read >= LIST_CAP ? undefined : read
  },
})

/** One page of runs, for a caller that pages nothing — a pinned chart. */
export const runsListQueryOptions = (
  filters: RunFilters & { limit?: number },
) => ({
  queryKey: [...runKeys.all, "pinned", filters] as const,
  queryFn: () => RunsService.readRuns(filters),
})

export const runOverviewQueryOptions = () => ({
  queryKey: runKeys.overview,
  queryFn: () => RunsService.readOverview(),
  refetchInterval: 30_000,
})

export const runQueryOptions = (id: string) => ({
  queryKey: runKeys.detail(id),
  queryFn: () => RunsService.readRun({ runId: id }),
  // A finished run never changes again, so only a live one is polled. The
  // socket's run_finished lands the last transition either way; this covers
  // the metrics and node rows filling in while it runs.
  refetchInterval: (query: { state: { data?: { status: string } } }) =>
    query.state.data && isLive(query.state.data.status)
      ? 5_000
      : (false as const),
})

/** One run's series, or every one of them when `name` is empty. */
export const runMetricsQueryOptions = (id: string, name = "") => ({
  queryKey: runKeys.metrics(id, name),
  queryFn: () => RunsService.readMetrics({ runId: id, name }),
})

/** What a comparison can be plotted against, beside another metric's name. */
export const STEP_AXIS = "step"
export const TIME_AXIS = "time"

export const compareQueryOptions = (
  ids: string[],
  metric: string,
  x = STEP_AXIS,
  refetchInterval?: number,
) => ({
  queryKey: runKeys.compare(ids, metric, x),
  queryFn: () => RunsService.compareMetric({ ids: ids.join(","), metric, x }),
  enabled: ids.length > 0 && Boolean(metric),
  ...(refetchInterval ? { refetchInterval } : {}),
})

/**
 * How many runs one selection carries.
 *
 * The selection is the address, which is what makes a comparison a link — and
 * a URL is not a place to put five hundred 22-character ids. A chart draws
 * five of them anyway; the rest of the room is for changing your mind.
 */
export const MAX_SELECTION = 50

export function useCancelRun() {
  const client = useQueryClient()
  return useMutation({
    mutationFn: (runId: string) => RunsService.cancelRun({ runId }),
    onSuccess: () => client.invalidateQueries({ queryKey: runKeys.all }),
  })
}

/**
 * Run the same thing again, as a run of its own.
 *
 * The engine keeps the flow, the inputs, the seed and the group, so a sweep
 * missing one config is completed rather than reissued — and lands on the new
 * run, since that is what there is to watch.
 */
export function useRetryRun() {
  const client = useQueryClient()
  const navigate = useNavigate()
  return useMutation({
    mutationFn: (runId: string) => RunsService.retryRun({ runId }),
    onSuccess: (run: { id: string }) => {
      client.invalidateQueries({ queryKey: runKeys.all })
      navigate({ to: "/runs/$id", params: { id: run.id } })
    },
  })
}

/**
 * Save an artifact to disk.
 *
 * Not a plain link: `/artifacts/{digest}` takes a bearer token, which an
 * anchor cannot carry. The bytes come through fetch and leave as an object
 * URL — the same trip the browser would have made, with the header on it.
 */
export async function downloadArtifact(digest: string, name: string) {
  const token = apiToken()
  const answer = await fetch(`${OpenAPI.BASE}/api/v1/artifacts/${digest}`, {
    headers: token ? { Authorization: `Bearer ${token}` } : {},
  })
  if (!answer.ok) throw new Error(`Could not read ${name}`)
  const url = URL.createObjectURL(await answer.blob())
  const link = document.createElement("a")
  link.href = url
  link.download = name
  link.click()
  URL.revokeObjectURL(url)
}

/**
 * What a flow declares it can be given, by input name.
 *
 * A batch flow's inputs *are* its parameters — a run supplies values for the
 * ones it names and takes the flow's own for the rest — so this is what turns
 * `run.params` from "what was passed" into "what the run actually ran with".
 *
 * The declarations are the flow's *current* ones, while a run carries the
 * `flow_version` it was submitted against. An input added since is shown on an
 * older run as a default it never actually received.
 */
export function useFlowInputs(flow: string | undefined) {
  const { data } = useQuery({
    ...flowQueryOptions(flow ?? ""),
    enabled: Boolean(flow),
  })
  const inputs = data?.definition.inputs
  return useMemo(
    () =>
      new Map<string, unknown>(
        (inputs ?? [])
          .filter((one) => Boolean(one.spec.name))
          .map((one) => [one.spec.name ?? "", one.initial ?? null]),
      ),
    [inputs],
  )
}

/**
 * Save the current selection as a file.
 *
 * The same trip `downloadArtifact` makes and for the same reason — the export
 * routes take a bearer token, which an anchor cannot carry. Not the generated
 * SDK either: it parses every body as JSON, and these stream csv.
 */
async function exportAs(what: "runs" | "metrics", query: URLSearchParams) {
  const token = apiToken()
  const answer = await fetch(
    `${OpenAPI.BASE}/api/v1/runs/export/${what}?${query}`,
    { headers: token ? { Authorization: `Bearer ${token}` } : {} },
  )
  if (!answer.ok) throw new Error(`Could not export the ${what}`)
  const url = URL.createObjectURL(await answer.blob())
  const link = document.createElement("a")
  link.href = url
  link.download = `${what}.${query.get("format") ?? "csv"}`
  link.click()
  URL.revokeObjectURL(url)
}

/** The runs themselves: one row each, with the parameters that varied. */
export const exportRuns = (query: URLSearchParams) => exportAs("runs", query)

/** Every recorded number of the selection, one row per point. */
export const exportMetrics = (query: URLSearchParams) =>
  exportAs("metrics", query)

/** How long a cancelled run is given to actually stop before delete gives up. */
const SETTLE_TRIES = 30
const SETTLE_WAIT_MS = 500

/**
 * Delete a run, cancelling it first if it is still going.
 *
 * The route refuses a live run rather than racing its driver, so the two steps
 * are the caller's to sequence. The wait is bounded and throws when it runs
 * out, which is what puts a stuck run in the partial-success toast by name
 * instead of hanging the button.
 *
 * ponytail: polling, because nothing pushes a run's status to a caller that is
 * not rendering it. The socket already carries run_finished if this ever needs
 * to be immediate.
 */
export async function cancelThenDelete(runId: string) {
  const run = await RunsService.readRun({ runId })
  if (isLive(run.status)) {
    await RunsService.cancelRun({ runId })
    let settled = false
    for (let tries = 0; tries < SETTLE_TRIES && !settled; tries++) {
      await new Promise((wake) => setTimeout(wake, SETTLE_WAIT_MS))
      settled = !isLive((await RunsService.readRun({ runId })).status)
    }
    if (!settled) throw new Error(`${shortId(runId)} did not stop`)
  }
  await RunsService.deleteRun({ runId })
}

/**
 * Why a finished run can have nothing to draw.
 *
 * A cache hit replays no emissions, so a restored node's curve is read back
 * from the run that recorded it. Deleting that run — deleting its flow does —
 * takes the curve with it, and then a reused run has a result and nothing to
 * draw. Said out loud rather than left as an empty chart, which reads as a fault.
 */
export const NO_CURVE =
  "No curve was recorded. A node restored from the cache is read back from the run that produced it, so this draws nothing once that run has been deleted — its outputs are still on the result."

/** A run id, short enough for a table cell. The tail is the random half. */
export const shortId = (id: string) => id.slice(-8)

/**
 * A series label from the comparison endpoint, cut down to fit a legend.
 *
 * It leads with the whole run id, which is 22 characters of mostly timestamp
 * beside four others.
 */
export const shortenRunLabel = (label: string) =>
  label.replace(/^\S+/, (id) => shortId(id))

/** A commit, at the length everyone reads one at. */
export const shortCommit = (commit: string) => commit.slice(0, 7)

/**
 * The parameter keys that differ across these runs.
 *
 * What makes a sweep readable: fifty runs of one flow share everything but the
 * two knobs that were swept, and those two are the only columns worth drawing.
 */
export function varyingKeys(runs: { params: Record<string, unknown> }[]) {
  if (runs.length < 2) return []
  const keys = new Set<string>()
  for (const run of runs)
    for (const key of Object.keys(run.params)) keys.add(key)
  return [...keys].filter((key) => {
    const first = JSON.stringify(runs[0].params[key])
    return runs.some((run) => JSON.stringify(run.params[key]) !== first)
  })
}

/** A parameter value, as narrow as it can be written. */
export function paramText(value: unknown): string {
  if (value === null || value === undefined) return "—"
  if (typeof value === "number" || typeof value === "boolean")
    return String(value)
  if (typeof value === "string") return value
  return JSON.stringify(value)
}

/** A run's parameters on one line, for a table cell. */
export const paramsSummary = (params: Record<string, unknown>) =>
  Object.entries(params)
    .map(([key, value]) => `${key}=${paramText(value)}`)
    .join("  ")