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

import { useQueryClient } from "@tanstack/react-query"
import { useEffect, useRef } from "react"

import { OpenAPI } from "@/client"
import { type LogLine, liveStore, type ValueSource } from "./liveStore"
import { flowKeys } from "./queries"

const RECONNECT_MIN = 1000

/** How many components want the socket open. */
let mounted = 0
const RECONNECT_MAX = 30000

type FlowEvent =
  | {
      type: "snapshot"
      values: Record<string, { value: unknown; ts: number | null }>
      nodes: { id: string; status: string; error?: string | null }[]
      paused?: string[]
      logs?: LogLine[]
    }
  | {
      type: "message_value"
      name: string
      value: unknown
      ts: number
      source?: ValueSource
    }
  | { type: "node_executed"; node: string; outputs: number }
  | { type: "node_error"; node: string; error: string }
  | { type: "node_status"; node: string; status: string; error?: string | null }
  | ({ type: "node_log" } & LogLine)
  | { type: "flow_paused"; flow: string; paused: boolean }
  | {
      type: "pipeline_rebuilt"
      nodes: { id: string; status: string; error?: string | null }[]
      paused?: string[]
    }

function socketUrl(): string {
  const base = String(OpenAPI.BASE || window.location.origin)
  const url = new URL("/api/v1/flows/ws", base)
  url.protocol = url.protocol === "https:" ? "wss:" : "ws:"
  // Browsers cannot set headers on a websocket handshake, so the token rides
  // in the query string.
  url.searchParams.set("token", localStorage.getItem("access_token") ?? "")
  return url.toString()
}

/**
 * Keeps one socket open for the editor, feeding the live store.
 *
 * @param onAuthFailure called when the server rejects the token, so the caller
 *   can send the user back to the login screen.
 */
export function useFlowSocket(onAuthFailure?: () => void): void {
  const socket = useRef<WebSocket | null>(null)
  const retry = useRef(RECONNECT_MIN)
  const timer = useRef<ReturnType<typeof setTimeout> | null>(null)
  const closed = useRef(false)
  const queryClient = useQueryClient()

  useEffect(() => {
    // The editor and a dashboard can both be mounted; one socket serves both,
    // and the second caller just rides along.
    mounted += 1
    if (mounted > 1) {
      return () => {
        mounted -= 1
      }
    }
    closed.current = false

    const connect = () => {
      if (closed.current) return
      const ws = new WebSocket(socketUrl())
      socket.current = ws

      ws.onopen = () => {
        retry.current = RECONNECT_MIN
        liveStore.setConnected(true)
      }

      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 ?? [])
            break
          case "message_value":
            liveStore.setValue(message.name, {
              value: message.value,
              ts: message.ts,
              source: message.source,
            })
            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,
            })
            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.
            queryClient.invalidateQueries({ queryKey: flowKeys.all })
            break
        }
      }

      ws.onclose = (event) => {
        liveStore.setConnected(false)
        if (closed.current) return
        if (event.code === 1008) {
          onAuthFailure?.()
          return
        }
        timer.current = setTimeout(connect, retry.current)
        retry.current = Math.min(retry.current * 2, RECONNECT_MAX)
      }
    }

    connect()

    return () => {
      mounted -= 1
      closed.current = true
      if (timer.current) clearTimeout(timer.current)
      socket.current?.close()
      liveStore.setConnected(false)
    }
  }, [onAuthFailure, queryClient])
}