import { useEffect, useRef, useCallback } from "react";
import type { ServerMessage, ClientMessage } from "../types";

export function useWebSocket(
  mapId: string,
  onMessage: (msg: ServerMessage) => void,
) {
  const wsRef = useRef<WebSocket | null>(null);
  // Keep onMessage stable via a ref to avoid reconnecting on every render
  const onMessageRef = useRef(onMessage);
  onMessageRef.current = onMessage;

  useEffect(() => {
    const proto = window.location.protocol === "https:" ? "wss:" : "ws:";
    // The browser automatically sends the siren_session cookie with the
    // WebSocket upgrade request — no manual token query param needed.
    const url = `${proto}//${window.location.host}/api/grid/maps/${mapId}/ws`;

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

    ws.onopen = () => {
      console.log(`[WS] Connected to map ${mapId}`);
    };

    ws.onmessage = (event: MessageEvent) => {
      try {
        const msg: ServerMessage = JSON.parse(event.data as string);
        onMessageRef.current(msg);
      } catch (err) {
        console.error("[WS] Failed to parse message:", err);
      }
    };

    ws.onerror = (err) => {
      console.error("[WS] Error:", err);
    };

    ws.onclose = () => {
      console.log(`[WS] Disconnected from map ${mapId}`);
    };

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

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

  return { send };
}
