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);
  /** A deferred teardown, cancelled when the effect re-runs — see the cleanup below. */
  const closeTimerRef = useRef<number | null>(null);
  const onMessageRef = useRef(onMessage);
  onMessageRef.current = onMessage;
  const onOpenRef = useRef(onOpen);
  onOpenRef.current = onOpen;

  /**
   * Messages typed before the socket was ready.
   *
   * `send` used to drop them: `readyState !== OPEN` returned, silently, with no error and no retry — so
   * pressing enter did nothing and the turn never happened. That window is not rare. React's dev
   * StrictMode double-invokes effects, so every socket is created, closed and recreated on mount, and a
   * reconnect after a drop reopens it again; with several chat panes on screen there are several sockets
   * doing this at once. One of them is always briefly not OPEN.
   *
   * Queued and flushed on open, in order. The mobile chat app does exactly this and for exactly this
   * reason — the composer is allowed to fire before the transport is ready.
   */
  const pendingRef = useRef<string[]>([]);

  const connect = () => {
    if (isCleaningUpRef.current) return;
    // CONNECTING counts as ours, not just OPEN. The deferred teardown below keeps a remount's socket
    // alive mid-handshake, and this is what reclaims it: checking only OPEN meant an effect re-run
    // built a SECOND socket and orphaned the first, which then stayed open forever with its own `open`
    // handler bailing on the socketRef mismatch.
    //
    // This is per-instance and does NOT explain the two sockets a /chat/new load opens — measured with
    // a WebSocket-constructor counter, those come from two separate `useChat` instances mounting, each
    // with its own refs. Unresolved, and tracked separately; both connect, so it reads as healthy.
    const existing = socketRef.current;
    if (existing && (existing.readyState === WebSocket.OPEN || existing.readyState === WebSocket.CONNECTING)) return;

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

    socket.addEventListener('open', () => {
      if (socketRef.current !== socket) return;
      setIsConnected(true);
      retryRef.current = 0;
      // BEFORE onOpen, deliberately: onOpen sends the resume/attach handshake, and anything the user
      // typed while connecting belongs after that, not in front of it.
      const queued = pendingRef.current;
      pendingRef.current = [];
      onOpenRef.current?.();
      for (const message of queued) socket.send(message);
    });

    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(() => {
    // A pending teardown from a remount that is about to be undone — see below.
    if (closeTimerRef.current !== null) {
      clearTimeout(closeTimerRef.current);
      closeTimerRef.current = null;
    }
    isCleaningUpRef.current = false;
    connect();

    return () => {
      /**
       * Close LATER, not now.
       *
       * Closing here directly is correct for a real unmount and disastrous for a remount, and this hook
       * cannot tell them apart at the moment it runs. React's dev StrictMode double-invokes every effect
       * (mount → unmount → mount), and a subtree that is re-created — a resolved transcript, a parent key
       * change — does the same. Each time, the socket was closed while still CONNECTING, the browser
       * logged "closed before the connection is established", and the replacement was closed in turn, so
       * the view could churn forever and never hold a connection while sitting on Disconnected.
       *
       * Deferring by a tick makes the two distinguishable. A remount re-runs the effect immediately and
       * cancels this timer, so the live socket is kept and the handshake completes. A real unmount has
       * nobody to cancel it and the socket closes a frame later, which costs nothing.
       */
      isCleaningUpRef.current = true;
      if (retryTimeoutRef.current !== null) {
        clearTimeout(retryTimeoutRef.current);
        retryTimeoutRef.current = null;
      }
      const socket = socketRef.current;
      closeTimerRef.current = window.setTimeout(() => {
        closeTimerRef.current = null;
        if (!isCleaningUpRef.current) return; // remounted: the effect above already reclaimed it
        if (socket) socket.close();
        if (socketRef.current === socket) socketRef.current = null;
      }, 0);
    };
  }, [url]);

  const send = (data: Record<string, unknown>) => {
    const socket = socketRef.current;
    const message = JSON.stringify(data);
    if (socket && socket.readyState === WebSocket.OPEN) {
      socket.send(message);
      return;
    }
    // Not open yet, or reconnecting. Hold it rather than dropping it — see `pendingRef`. Bounded so a
    // socket that never comes back cannot grow this without limit; the oldest go first, because the
    // newest message is the one the user is still waiting on.
    pendingRef.current.push(message);
    if (pendingRef.current.length > 50) pendingRef.current.shift();
  };

  return { isConnected, send };
};
