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

export type WebSocketStatus =
  | "connecting"
  | "connected"
  | "disconnected"
  | "error";

interface WebSocketOptions {
  url: string;
  onMessage?: (data: unknown) => void;
  onConnect?: () => void;
  onDisconnect?: () => void;
  onError?: (error: Event) => void;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  heartbeatInterval?: number;
}

interface WebSocketState {
  status: WebSocketStatus;
  reconnectAttempts: number;
}

// Exponential backoff delays: 1s, 2s, 4s, 8s, 16s, 30s (max)
const RECONNECT_DELAYS = [1000, 2000, 4000, 8000, 16000, 30000];

export function useWebSocket({
  url,
  onMessage,
  onConnect,
  onDisconnect,
  onError,
  maxReconnectAttempts = 10,
  heartbeatInterval = 30000,
}: WebSocketOptions) {
  const [state, setState] = useState<WebSocketState>({
    status: "disconnected",
    reconnectAttempts: 0,
  });

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(
    null,
  );
  const heartbeatIntervalRef = useRef<ReturnType<typeof setInterval> | null>(
    null,
  );
  const reconnectAttemptsRef = useRef(0);
  const isManualCloseRef = useRef(false);
  const subscriptionsRef = useRef<Set<string>>(new Set());

  // Use refs for callbacks to prevent unnecessary reconnection loops
  const onMessageRef = useRef(onMessage);
  const onConnectRef = useRef(onConnect);
  const onDisconnectRef = useRef(onDisconnect);
  const onErrorRef = useRef(onError);

  useEffect(() => {
    onMessageRef.current = onMessage;
    onConnectRef.current = onConnect;
    onDisconnectRef.current = onDisconnect;
    onErrorRef.current = onError;
  }, [onMessage, onConnect, onDisconnect, onError]);

  const getReconnectDelay = useCallback((attempt: number): number => {
    const index = Math.min(attempt, RECONNECT_DELAYS.length - 1);
    return RECONNECT_DELAYS[index];
  }, []);

  const clearReconnectTimeout = useCallback(() => {
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }
  }, []);

  const clearHeartbeatInterval = useCallback(() => {
    if (heartbeatIntervalRef.current) {
      clearInterval(heartbeatIntervalRef.current);
      heartbeatIntervalRef.current = null;
    }
  }, []);

  const startHeartbeat = useCallback(() => {
    clearHeartbeatInterval();
    heartbeatIntervalRef.current = setInterval(() => {
      if (wsRef.current?.readyState === WebSocket.OPEN) {
        wsRef.current.send(JSON.stringify({ type: "ping" }));
      }
    }, heartbeatInterval);
  }, [heartbeatInterval, clearHeartbeatInterval]);

  const subscribe = useCallback((channel: string) => {
    subscriptionsRef.current.add(channel);
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: "subscribe", channel }));
    }
  }, []);

  const unsubscribe = useCallback((channel: string) => {
    subscriptionsRef.current.delete(channel);
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: "unsubscribe", channel }));
    }
  }, []);

  const resubscribeAll = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      subscriptionsRef.current.forEach((channel) => {
        wsRef.current?.send(JSON.stringify({ type: "subscribe", channel }));
      });
    }
  }, []);

  const connect = useCallback(() => {
    if (
      wsRef.current?.readyState === WebSocket.OPEN ||
      wsRef.current?.readyState === WebSocket.CONNECTING
    ) {
      return;
    }

    setState((prev) => ({ ...prev, status: "connecting" }));
    isManualCloseRef.current = false;

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

      ws.onopen = () => {
        reconnectAttemptsRef.current = 0;
        setState({
          status: "connected",
          reconnectAttempts: 0,
        });
        onConnectRef.current?.();
        startHeartbeat();
        // Resubscribe to channels after reconnection
        resubscribeAll();
      };

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          // Ignore pong responses
          if (data.type === "pong") return;
          onMessageRef.current?.(data);
        } catch {
          // If not JSON, pass raw data
          onMessageRef.current?.(event.data);
        }
      };

      ws.onclose = () => {
        setState((prev) => ({
          ...prev,
          status: "disconnected",
        }));
        clearHeartbeatInterval();
        onDisconnectRef.current?.();

        // Auto reconnect if not manually closed
        if (
          !isManualCloseRef.current &&
          reconnectAttemptsRef.current < maxReconnectAttempts
        ) {
          reconnectAttemptsRef.current += 1;
          const delay = getReconnectDelay(reconnectAttemptsRef.current);
          setState((prev) => ({
            ...prev,
            reconnectAttempts: reconnectAttemptsRef.current,
          }));
          reconnectTimeoutRef.current = setTimeout(() => {
            connect();
          }, delay);
        }
      };

      ws.onerror = (error) => {
        setState((prev) => ({ ...prev, status: "error" }));
        // Only log errors after first reconnect attempt to avoid transient errors
        if (reconnectAttemptsRef.current > 0) {
          console.error(`[WebSocket] Connection error to ${url}:`, error);
        }
        onErrorRef.current?.(error);
      };
    } catch (error) {
      setState((prev) => ({ ...prev, status: "error" }));
      console.error(`[WebSocket] Initialization error:`, error);
      onErrorRef.current?.(error as Event);
    }
  }, [
    url,
    maxReconnectAttempts,
    startHeartbeat,
    clearHeartbeatInterval,
    getReconnectDelay,
    resubscribeAll,
  ]);

  const disconnect = useCallback(() => {
    isManualCloseRef.current = true;
    clearReconnectTimeout();
    clearHeartbeatInterval();
    subscriptionsRef.current.clear();
    if (wsRef.current) {
      wsRef.current.close();
      wsRef.current = null;
    }
    setState({
      status: "disconnected",
      reconnectAttempts: 0,
    });
  }, [clearReconnectTimeout, clearHeartbeatInterval]);

  const sendMessage = useCallback((data: unknown) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(
        typeof data === "string" ? data : JSON.stringify(data),
      );
      return true;
    }
    return false;
  }, []);

  useEffect(() => {
    connect();

    return () => {
      disconnect();
    };
  }, [connect, disconnect]);

  return {
    ...state,
    isConnected: state.status === "connected",
    isConnecting: state.status === "connecting",
    connect,
    disconnect,
    sendMessage,
    subscribe,
    unsubscribe,
  };
}
