import { useState, useEffect, useRef } from 'react';

type UseChatWebSocketParams = {
  url: string;
  onMessage: (data: unknown) => void;
  // Called on every (re)connect once the socket is OPEN — used to send a resume-cursor so a reconnect
  // re-binds to the session and replays missed events (no manual navigate-away/back).
  onOpen?: () => void;
};

export const useChatWebSocket = ({ url, onMessage, onOpen }: UseChatWebSocketParams) => {
  const [isConnected, setIsConnected] = useState(false);

  const socketRef = useRef<WebSocket | null>(null);
  const retryRef = useRef(0);
  const retryTimeoutRef = useRef<number | null>(null);
  const isCleaningUpRef = useRef(false);
  const onMessageRef = useRef(onMessage);
  onMessageRef.current = onMessage;
  const onOpenRef = useRef(onOpen);
  onOpenRef.current = onOpen;

  const connect = () => {
    if (isCleaningUpRef.current) return;
    if (socketRef.current && socketRef.current.readyState === WebSocket.OPEN) return;

    const socket = new WebSocket(url);
    socketRef.current = socket;

    socket.addEventListener('open', () => {
      if (socketRef.current !== socket) return;
      setIsConnected(true);
      retryRef.current = 0;
      onOpenRef.current?.();
    });

    socket.addEventListener('message', (ev) => {
      try {
        const data = JSON.parse(typeof ev.data === 'string' ? ev.data : new TextDecoder().decode(ev.data));
        onMessageRef.current(data);
      } catch {
        // ignore unparseable messages
      }
    });

    socket.addEventListener('close', () => {
      if (isCleaningUpRef.current) return;
      if (socketRef.current !== socket) return;
      setIsConnected(false);
      const retry = (retryRef.current ?? 0) + 1;
      retryRef.current = retry;
      const delay = Math.min(5000, 300 * retry);
      retryTimeoutRef.current = window.setTimeout(connect, delay);
    });

    socket.addEventListener('error', () => {
      socket.close();
    });
  };

  useEffect(() => {
    isCleaningUpRef.current = false;
    connect();
    return () => {
      isCleaningUpRef.current = true;
      if (retryTimeoutRef.current !== null) {
        clearTimeout(retryTimeoutRef.current);
        retryTimeoutRef.current = null;
      }
      if (socketRef.current) {
        socketRef.current.close();
        socketRef.current = null;
      }
    };
  }, [url]);

  const send = (data: Record<string, unknown>) => {
    const socket = socketRef.current;
    if (!socket || socket.readyState !== WebSocket.OPEN) return;
    socket.send(JSON.stringify(data));
  };

  return { isConnected, send };
};
