frontend/src/lib/useChatSession.ts

/**
 * useChatSession — single hook that owns the WebSocket connection and all derived state.
 *
 * Production behaviour:
 *   - On mount, opens /ws (with ?session_id=<persisted> if we have one in localStorage).
 *   - Receives the `session` envelope, persists the session_id so a refresh resumes.
 *   - Streams trace events live (they land in `traceEvents` before the reply lands).
 *   - On reconnect (server bounce / network blip) we keep the same session_id so the
 *     backend's AgentLoop is preserved.
 *   - Exposes minimal imperative methods: send(), reset(), reconnect().
 */
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import type {
  AgentStateSnapshot,
  ChatMessage,
  ClientMessage,
  ConnectionStatus,
  ServerMessage,
  TraceEvent,
  UiHint,
} from "../types";

const SESSION_KEY = "ramco-chat.session_id";
const WS_URL = `${location.protocol === "https:" ? "wss" : "ws"}://${location.host}/ws`;
const RECONNECT_DELAY_MS = 1500;
const MAX_RECONNECT_DELAY_MS = 10_000;

function uid(): string {
  return Math.random().toString(36).slice(2, 10);
}

export type ChatSession = {
  status: ConnectionStatus;
  sessionId: string | null;
  resumed: boolean;
  availableJourneys: string[];
  state: AgentStateSnapshot | null;
  messages: ChatMessage[];
  /** Trace events for the most recent turn (cleared at start of each new turn). */
  liveTrace: TraceEvent[];
  /** All trace events ever, for a "show full history" pane. */
  allTrace: TraceEvent[];
  /** True while the agent is processing the latest user message. */
  thinking: boolean;
  /** The latest assistant ui_hint, used by the smart-control renderer. */
  activeUiHint: UiHint | null;
  /** When the latest turn produced candidate_journeys (ambiguity), they live here. */
  candidateJourneys: string[];
  send(text: string): void;
  reset(): void;
  reconnect(): void;
};

export function useChatSession(): ChatSession {
  const [status, setStatus] = useState<ConnectionStatus>("connecting");
  const [sessionId, setSessionId] = useState<string | null>(() => localStorage.getItem(SESSION_KEY));
  const [resumed, setResumed] = useState(false);
  const [availableJourneys, setAvailableJourneys] = useState<string[]>([]);
  const [state, setState] = useState<AgentStateSnapshot | null>(null);
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [liveTrace, setLiveTrace] = useState<TraceEvent[]>([]);
  const [allTrace, setAllTrace] = useState<TraceEvent[]>([]);
  const [thinking, setThinking] = useState(false);
  const [activeUiHint, setActiveUiHint] = useState<UiHint | null>(null);
  const [candidateJourneys, setCandidateJourneys] = useState<string[]>([]);

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectAttemptRef = useRef(0);
  const reconnectTimerRef = useRef<number | null>(null);
  const intentionalCloseRef = useRef(false);

  const sendClient = useCallback((msg: ClientMessage) => {
    const ws = wsRef.current;
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(msg));
    }
  }, []);

  const connect = useCallback(() => {
    const stored = localStorage.getItem(SESSION_KEY);
    const url = stored ? `${WS_URL}?session_id=${encodeURIComponent(stored)}` : WS_URL;
    setStatus("connecting");
    intentionalCloseRef.current = false;

    const ws = new WebSocket(url);
    wsRef.current = ws;

    ws.onopen = () => {
      setStatus("open");
      reconnectAttemptRef.current = 0;
    };

    ws.onmessage = (ev) => {
      let msg: ServerMessage;
      try {
        msg = JSON.parse(ev.data);
      } catch {
        return;
      }
      handleServerMessage(msg);
    };

    ws.onclose = () => {
      setStatus("closed");
      wsRef.current = null;
      setThinking(false);
      if (!intentionalCloseRef.current) {
        scheduleReconnect();
      }
    };

    ws.onerror = () => {
      setStatus("error");
    };
  }, []);

  const scheduleReconnect = useCallback(() => {
    if (reconnectTimerRef.current !== null) return;
    const attempt = ++reconnectAttemptRef.current;
    const delay = Math.min(RECONNECT_DELAY_MS * attempt, MAX_RECONNECT_DELAY_MS);
    reconnectTimerRef.current = window.setTimeout(() => {
      reconnectTimerRef.current = null;
      connect();
    }, delay);
  }, [connect]);

  const handleServerMessage = useCallback((msg: ServerMessage) => {
    switch (msg.type) {
      case "session": {
        setSessionId(msg.session_id);
        localStorage.setItem(SESSION_KEY, msg.session_id);
        setResumed(msg.resumed);
        setAvailableJourneys(msg.available_journeys);
        setState(msg.state);
        // If the server reports a non-trivial resumed state, surface a system note.
        if (msg.resumed && msg.state.turn_count > 0) {
          setMessages((m) => [
            ...m,
            {
              id: uid(),
              role: "system",
              text: `Resumed your session — ${msg.state.turn_count} turns already on the record.`,
              ts: Date.now(),
            },
          ]);
        }
        break;
      }
      case "turn_start": {
        setThinking(true);
        setLiveTrace([]);
        setActiveUiHint(null);
        setCandidateJourneys([]);
        // We optimistically appended the user message in send(); nothing to do here.
        break;
      }
      case "trace": {
        setLiveTrace((t) => [...t, msg.event]);
        setAllTrace((t) => [...t, msg.event]);
        break;
      }
      case "turn_complete": {
        setThinking(false);
        setState(msg.state);
        setActiveUiHint(msg.ui_hint);
        // Extract candidate_journeys from the most recent clarify_required trace event.
        const candidates = lastCandidateJourneysFromTrace(allTraceRef.current.concat(liveTraceRef.current));
        setCandidateJourneys(candidates);
        setMessages((m) => [
          ...m,
          {
            id: uid(),
            role: "assistant",
            text: msg.reply,
            ts: Date.now(),
            ui_hint: msg.ui_hint,
            candidates,
            terminated: msg.terminated,
          },
        ]);
        break;
      }
      case "error": {
        setThinking(false);
        setMessages((m) => [
          ...m,
          { id: uid(), role: "system", text: `Error: ${msg.message}`, ts: Date.now(), level: "error" },
        ]);
        break;
      }
      case "pong":
        break;
    }
  }, []);

  // refs so handleServerMessage can read current trace state inside its closure
  const allTraceRef = useRef<TraceEvent[]>([]);
  const liveTraceRef = useRef<TraceEvent[]>([]);
  useEffect(() => { allTraceRef.current = allTrace; }, [allTrace]);
  useEffect(() => { liveTraceRef.current = liveTrace; }, [liveTrace]);

  useEffect(() => {
    connect();
    return () => {
      intentionalCloseRef.current = true;
      if (reconnectTimerRef.current !== null) clearTimeout(reconnectTimerRef.current);
      wsRef.current?.close();
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);

  // Keep-alive ping every 30s so dev proxies don't drop idle sockets.
  useEffect(() => {
    const id = window.setInterval(() => sendClient({ type: "ping" }), 30_000);
    return () => clearInterval(id);
  }, [sendClient]);

  const send = useCallback(
    (text: string) => {
      const trimmed = text.trim();
      if (!trimmed) return;
      // Optimistically render the user bubble; server will echo via turn_start.
      setMessages((m) => [...m, { id: uid(), role: "user", text: trimmed, ts: Date.now() }]);
      setLiveTrace([]);
      setActiveUiHint(null);
      setCandidateJourneys([]);
      sendClient({ type: "user_message", text: trimmed });
    },
    [sendClient],
  );

  const reset = useCallback(() => {
    sendClient({ type: "reset" });
    setMessages([]);
    setAllTrace([]);
    setLiveTrace([]);
    setActiveUiHint(null);
    setCandidateJourneys([]);
  }, [sendClient]);

  const reconnect = useCallback(() => {
    intentionalCloseRef.current = true;
    wsRef.current?.close();
    intentionalCloseRef.current = false;
    setTimeout(connect, 50);
  }, [connect]);

  return useMemo(
    () => ({
      status,
      sessionId,
      resumed,
      availableJourneys,
      state,
      messages,
      liveTrace,
      allTrace,
      thinking,
      activeUiHint,
      candidateJourneys,
      send,
      reset,
      reconnect,
    }),
    [
      status,
      sessionId,
      resumed,
      availableJourneys,
      state,
      messages,
      liveTrace,
      allTrace,
      thinking,
      activeUiHint,
      candidateJourneys,
      send,
      reset,
      reconnect,
    ],
  );
}

function lastCandidateJourneysFromTrace(events: TraceEvent[]): string[] {
  for (let i = events.length - 1; i >= 0; i--) {
    const e = events[i];
    if (e.kind === "clarify_required" && Array.isArray((e.data as any)?.candidates)) {
      return (e.data as any).candidates as string[];
    }
    if (e.kind === "turn_begin") break; // only look within the current turn
  }
  return [];
}