"use client";

import { useEffect, useRef, useState } from "react";
import { API_BASE_URL, getAccessToken } from "./api";
import type { RunStatus, Step } from "./types";

const TERMINAL = new Set(["succeeded", "failed", "cancelled", "waiting_approval"]);

export interface RunStreamState {
  steps: Step[];
  status: RunStatus | null;
  connected: boolean;
}

/**
 * Streams a run over WebSocket: the server replays persisted steps, then pushes
 * live `step` / `run.status` events until the run is terminal.
 */
export function useRunStream(runId: string): RunStreamState {
  const [steps, setSteps] = useState<Record<string, Step>>({});
  const [status, setStatus] = useState<RunStatus | null>(null);
  const [connected, setConnected] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    setSteps({});
    setStatus(null);
    const token = getAccessToken();
    if (!token) return;
    const wsBase = API_BASE_URL.replace(/^http/, "ws");
    const ws = new WebSocket(`${wsBase}/runs/${runId}/stream?token=${encodeURIComponent(token)}`);
    wsRef.current = ws;

    ws.onopen = () => setConnected(true);
    ws.onclose = () => setConnected(false);
    ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data);
        if (msg.type === "step" && msg.step) {
          setSteps((cur) => ({ ...cur, [msg.step.id]: msg.step }));
        } else if (msg.type === "run.status") {
          setStatus(msg.status);
          if (TERMINAL.has(msg.status)) ws.close();
        }
      } catch {
        /* ignore malformed frames */
      }
    };

    return () => {
      wsRef.current = null;
      ws.close();
    };
  }, [runId]);

  const ordered = Object.values(steps).sort((a, b) => a.idx - b.idx);
  return { steps: ordered, status, connected };
}
