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

export interface PriceUpdate {
  id: string;
  yes: number;
  no: number;
  vol: number;
}

interface WsMessage {
  type: 'prices';
  data: PriceUpdate[];
  ts: string;
}

interface WsState {
  prices: Map<string, PriceUpdate>;
  connected: boolean;
  /** market IDs that changed in the last update (for flash animation) */
  changed: Set<string>;
}

const RECONNECT_BASE = 1000;
const RECONNECT_MAX = 30000;

export function useWebSocket() {
  const [state, setState] = useState<WsState>({
    prices: new Map(),
    connected: false,
    changed: new Set(),
  });

  const wsRef = useRef<WebSocket | null>(null);
  const retriesRef = useRef<number>(0);
  const timerRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
  const mountedRef = useRef<boolean>(true);
  const connectRef = useRef<(() => void) | undefined>(undefined);

  const scheduleReconnect = useCallback(() => {
    const delay = Math.min(RECONNECT_BASE * 2 ** retriesRef.current, RECONNECT_MAX);
    retriesRef.current++;
    timerRef.current = setTimeout(() => {
      if (mountedRef.current) connectRef.current?.();
    }, delay);
  }, []);

  const connect = useCallback(() => {
    if (!mountedRef.current) return;

    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
    const url = `${proto}//${location.host}/ws/prices`;

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

    ws.onopen = () => {
      if (!mountedRef.current) return;
      retriesRef.current = 0;
      setState(prev => ({ ...prev, connected: true }));
    };

    ws.onmessage = (ev) => {
      if (!mountedRef.current) return;
      try {
        const msg: WsMessage = JSON.parse(ev.data);
        if (msg.type === 'prices' && Array.isArray(msg.data)) {
          setState(prev => {
            const next = new Map(prev.prices);
            const changed = new Set<string>();

            for (const p of msg.data) {
              const old = next.get(p.id);
              if (old && (old.yes !== p.yes || old.no !== p.no)) {
                changed.add(p.id);
              }
              next.set(p.id, p);
            }

            return { prices: next, connected: true, changed };
          });

          // Clear changed set after animation duration (600ms)
          setTimeout(() => {
            if (mountedRef.current) {
              setState(prev => ({ ...prev, changed: new Set() }));
            }
          }, 600);
        }
      } catch { /* ignore malformed */ }
    };

    ws.onclose = () => {
      if (!mountedRef.current) return;
      setState(prev => ({ ...prev, connected: false }));
      scheduleReconnect();
    };

    ws.onerror = () => {
      ws.close();
    };
  }, [scheduleReconnect]);

  useEffect(() => {
    connectRef.current = connect;
  }, [connect]);

  useEffect(() => {
    mountedRef.current = true;
    const t = setTimeout(connect, 0);

    return () => {
      mountedRef.current = false;
      clearTimeout(t);
      clearTimeout(timerRef.current);
      wsRef.current?.close();
    };
  }, [connect]);

  return state;
}