app/frontend/src/components/Flow/useFlowSocket.ts

import type { QueryClient } from "@tanstack/react-query"
import { useQueryClient } from "@tanstack/react-query"
import { useEffect } from "react"

import { OpenAPI } from "@/client"
import { dashboardKeys, panelKeys } from "@/components/Dashboard/queries"
import { connectionStore } from "@/lib/connectionStore"
import { apiToken } from "@/lib/portal"
import { type LogLine, liveStore, type ValueSource } from "./liveStore"
import { flowKeys } from "./queries"

const RECONNECT_MIN = 1000
const RECONNECT_MAX = 30000
/**
 * Consecutive token rejections before the session is given up on.
 *
 * A 1008 means the server would not take the token. Under a portal that is
 * routine — the handoff token lasts twelve hours against a local session's
 * eight days — so it has to be recoverable rather than fatal: the token is read
 * again for every attempt, so a refreshed one is picked up. A session that has
 * genuinely been revoked fails all three and falls through to the caller's auth
 * handler instead of retrying forever.
 */
const AUTH_ATTEMPTS = 3

type FlowEvent =
  | {
      type: "snapshot"
      values: Record<string, { value: unknown; ts: number | null }>
      nodes: { id: string; status: string; error?: string | null }[]
      paused?: string[]
      logs?: LogLine[]
      /** Emission counts per qualified node. Absent on older installations. */
      emits?: Record<string, number>
    }
  | {
      type: "message_value"
      name: string
      value: unknown
      ts: number
      source?: ValueSource
    }
  | { type: "node_started"; node: string }
  | {
      type: "node_executed"
      flow?: string
      node: string
      outputs: number
      duration_ms?: number
    }
  | {
      type: "node_error"
      flow?: string
      node: string
      error: string
      ts?: number
    }
  | { type: "node_status"; node: string; status: string; error?: string | null }
  | ({ type: "node_log" } & LogLine)
  | { type: "flow_paused"; flow: string; paused: boolean }
  | {
      type: "node_health"
      flow?: string
      node: string
      health: "ok" | "down" | "unknown"
      detail?: string | null
      ts?: number
    }
  | { type: "flow_quarantined"; flow: string; error?: string; ts?: number }
  | { type: "engine_degraded"; reason?: string; ts?: number }
  | { type: "engine_fatal"; reason?: string; ts?: number }
  | {
      type: "cascade_dropped"
      flow?: string
      node?: string
      deliveries?: number
      ts?: number
    }
  | {
      type: "queue_unavailable"
      flow?: string
      node?: string
      error?: string
      ts?: number
    }
  | {
      type: "pipeline_rebuilt"
      nodes: { id: string; status: string; error?: string | null }[]
      paused?: string[]
    }
  | { type: "dashboard_changed"; dashboard?: string; ts?: number }

function socketUrl(): string {
  const base = String(OpenAPI.BASE || window.location.origin)
  // Concatenated rather than resolved: an absolute path as the second argument
  // to `new URL` discards the base's own path, which under a portal
  // (`https://host/i/{id}`) would aim the socket at the wrong place entirely.
  const url = new URL(`${base.replace(/\/$/, "")}/api/v1/flows/ws`)
  url.protocol = url.protocol === "https:" ? "wss:" : "ws:"
  // Browsers cannot set headers on a websocket handshake, so the token rides
  // in the query string. Read per attempt, so a reconnect after the old one
  // expired carries whatever is current rather than what we started with.
  url.searchParams.set("token", apiToken())
  return url.toString()
}

/*
 * One socket for the page, owned here rather than by a component.
 *
 * Several places want live events at once — the app shell, the brain graph on
 * Home, the editor — and React mounts and unmounts them in an order none of
 * them controls: passive effects run children before parents, and StrictMode
 * runs mount → unmount → mount. Keeping the socket inside whichever hook
 * instance ran first made that one its owner, so leaving Home closed the socket
 * while the shell kept the count above zero and nothing ever reopened it. The
 * connection lives out here now with a plain refcount; a component only says
 * whether it is still watching.
 */
let socket: WebSocket | null = null
let timer: ReturnType<typeof setTimeout> | null = null
let retry = RECONNECT_MIN
/** How many mounted components want the socket open. */
let watchers = 0
/** Consecutive 1008s, cleared by a handshake the server accepts. */
let rejected = 0
let client: QueryClient | null = null
const authHandlers = new Set<() => void>()

function schedule() {
  if (timer) return
  timer = setTimeout(() => {
    timer = null
    connect()
  }, retry)
  retry = Math.min(retry * 2, RECONNECT_MAX)
}

function connect() {
  if (watchers === 0 || socket || timer) return
  const ws = new WebSocket(socketUrl())
  socket = ws

  ws.onopen = () => {
    retry = RECONNECT_MIN
    rejected = 0
    liveStore.setConnected(true)
    connectionStore.setSocketOpen(true)
    // Whatever happened while the socket was down was missed, so nothing
    // held in cache can be trusted to still be current.
    client?.invalidateQueries()
  }

  ws.onmessage = (event) => {
    const message: FlowEvent = JSON.parse(event.data)
    switch (message.type) {
      case "snapshot":
        liveStore.setValues(message.values)
        liveStore.setStatuses(message.nodes ?? [])
        liveStore.setPausedFlows(message.paused ?? [])
        liveStore.setLogs(message.logs ?? [])
        // Missing against an installation older than this bundle; the graph
        // then starts from zero the way it always did.
        liveStore.setEmits(message.emits ?? {})
        break
      case "message_value":
        liveStore.setValue(message.name, {
          value: message.value,
          ts: message.ts,
          source: message.source,
        })
        break
      case "node_started":
        liveStore.setStatus(message.node, { status: "running" })
        break
      case "node_executed":
        liveStore.setStatus(message.node, { status: "success" })
        if (message.outputs > 0) liveStore.recordEmit(message.node)
        break
      case "node_error":
        liveStore.setStatus(message.node, {
          status: "error",
          error: message.error,
        })
        liveStore.recordEngineEvent({
          type: message.type,
          flow: message.flow,
          node: message.node,
          detail: message.error,
          ts: message.ts ?? Date.now() / 1000,
        })
        break
      case "node_health":
        liveStore.setHealth(message.node, {
          health: message.health,
          detail: message.detail,
        })
        if (message.health === "down") {
          liveStore.recordEngineEvent({
            type: message.type,
            flow: message.flow,
            node: message.node,
            detail: message.detail ?? "Reported itself down.",
            ts: message.ts ?? Date.now() / 1000,
          })
        }
        break
      case "flow_quarantined":
      case "engine_degraded":
      case "engine_fatal":
      case "cascade_dropped":
      case "queue_unavailable":
        liveStore.recordEngineEvent({
          type: message.type,
          flow: "flow" in message ? message.flow : undefined,
          node: "node" in message ? message.node : undefined,
          detail:
            ("error" in message ? message.error : undefined) ??
            ("reason" in message ? message.reason : undefined) ??
            "",
          ts: message.ts ?? Date.now() / 1000,
        })
        break
      case "node_status":
        liveStore.setStatus(message.node, {
          status: message.status as "active" | "error",
          error: message.error,
        })
        break
      case "node_log":
        liveStore.appendLog(message)
        break
      case "flow_paused":
        liveStore.setPaused(message.flow, message.paused)
        break
      case "pipeline_rebuilt":
        liveStore.setStatuses(message.nodes)
        liveStore.setPausedFlows(message.paused ?? [])
        // Someone published or started a flow, here or in another tab: the
        // markers on the flow chips are stale until the list is refetched.
        client?.invalidateQueries({ queryKey: flowKeys.all })
        break
      case "dashboard_changed":
        // Someone published, or changed which dashboards a panel shows. Refetch
        // rather than reload: a wall screen must not blank or sign in again.
        // `exact` matters — the list key is a prefix of every detail key,
        // including the draft an editor may have open.
        client?.invalidateQueries({ queryKey: dashboardKeys.all, exact: true })
        client?.invalidateQueries({
          queryKey: message.dashboard
            ? dashboardKeys.detail(message.dashboard)
            : panelKeys.all,
        })
        break
    }
  }

  ws.onclose = (event) => {
    // A socket we already dropped: its close says nothing about the connection
    // we want now.
    if (socket !== ws) return
    socket = null
    liveStore.setConnected(false)
    connectionStore.setSocketOpen(false)
    if (watchers === 0) return
    if (event.code === 1008) {
      rejected += 1
      if (rejected >= AUTH_ATTEMPTS) {
        for (const handler of authHandlers) handler()
        return
      }
    }
    // 1013 is the portal saying the installation is not attached — the one
    // close code that means "offline" rather than "the socket dropped".
    if (event.code === 1013) {
      connectionStore.setOffline(null)
    }
    schedule()
  }
}

function release(onAuthFailure?: () => void) {
  watchers -= 1
  if (onAuthFailure) authHandlers.delete(onAuthFailure)
  if (watchers > 0) return
  if (timer) {
    clearTimeout(timer)
    timer = null
  }
  const ws = socket
  socket = null
  ws?.close()
  liveStore.setConnected(false)
  // Its `onclose` is dropped by the guard above, so say so here instead.
  connectionStore.setSocketOpen(false)
}

/**
 * Keeps one socket open while anything on the page wants live events.
 *
 * @param onAuthFailure called once the server has refused the token often
 *   enough to mean the session is gone, so the caller can send the user back to
 *   the login screen.
 */
export function useFlowSocket(onAuthFailure?: () => void): void {
  const queryClient = useQueryClient()

  useEffect(() => {
    client = queryClient
    if (onAuthFailure) authHandlers.add(onAuthFailure)
    watchers += 1
    connect()
    return () => release(onAuthFailure)
  }, [onAuthFailure, queryClient])
}