import React, {
  createContext,
  useContext,
  useEffect,
  useRef,
  useState,
  useCallback,
} from "react";
import { useAuth } from "./AuthContext";
import { getApiBaseUrl } from "./api";
import { logger } from "./utils/logger";

const WebSocketContext = createContext(null);

export const useWebSocket = () => {
  const context = useContext(WebSocketContext);
  if (!context) {
    throw new Error("useWebSocket must be used within a WebSocketProvider");
  }
  return context;
};

export const WebSocketProvider = ({ children }) => {
  const { user } = useAuth();
  const [lastMessage, setLastMessage] = useState(null);
  const [isConnected, setIsConnected] = useState(false);
  const [isFallback, setIsFallback] = useState(false);
  const ws = useRef(null);
  const reconnectTimeout = useRef(null);
  const pendingSubscriptions = useRef(new Set());
  const connectRef = useRef(null);
  const reconnectAttempts = useRef(0);
  const isConnecting = useRef(false);
  const MAX_RECONNECT_ATTEMPTS = 5;

  const connect = useCallback(async () => {
    // If not authenticated, don't attempt to connect
    if (!user) {
      return;
    }

    // If already connected, or explicitly connecting, don't reconnect
    if (
      ws.current &&
      (ws.current.readyState === WebSocket.OPEN ||
        ws.current.readyState === WebSocket.CONNECTING)
    ) {
      return;
    }

    if (isConnecting.current) {
      return;
    }

    if (typeof window === "undefined") return;

    // Check if we should fallback
    if (reconnectAttempts.current >= MAX_RECONNECT_ATTEMPTS) {
      logger.warn(
        "Max reconnect attempts reached. Switching to polling fallback.",
      );
      // Yield to avoid synchronous state update in effect
      setTimeout(() => {
        setIsFallback(true);
        setIsConnected(false);
      }, 0);
      return;
    }

    isConnecting.current = true;

    // Determine WebSocket URL based on API base URL or window location
    const baseUrl = getApiBaseUrl();
    let wsUrl;
    if (baseUrl) {
      // Replace http/https with ws/wss
      wsUrl = baseUrl.replace(/^http(s?):/, "ws$1:") + "/ws";
    } else {
      const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
      const host = window.location.host;
      wsUrl = `${protocol}//${host}/ws`;
    }

    // Append access token if available
    const token = localStorage.getItem("access_token");
    if (token) {
      wsUrl += `?token=${encodeURIComponent(token)}`;
    }

    // Rely on token in query param or cookies for authentication
    logger.debug(
      `Connecting to WebSocket at ${wsUrl.replace(/\?token=.*$/, "?token=***")}`,
    );

    try {
      const socket = new WebSocket(wsUrl);
      ws.current = socket;

      socket.onopen = () => {
        logger.debug("WebSocket Connected");
        setIsConnected(true);
        setIsFallback(false);
        reconnectAttempts.current = 0; // Reset attempts on success
        isConnecting.current = false;

        if (pendingSubscriptions.current.size > 0) {
          logger.debug(
            `Flushing ${pendingSubscriptions.current.size} pending subscriptions`,
          );
          pendingSubscriptions.current.forEach((topic) => {
            socket.send(JSON.stringify({ action: "subscribe", topic }));
          });
        }
      };

      socket.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          logger.debug(`[WebSocket Message] Topic: ${data.topic || "unknown"}`);
          setLastMessage(data);
        } catch (e) {
          logger.error("[WebSocket Error] Failed to parse message", e);
        }
      };

      socket.onclose = (event) => {
        logger.debug(
          `WebSocket Disconnected. Code: ${event.code}, Reason: ${event.reason}`,
        );
        setIsConnected(false);
        ws.current = null;
        isConnecting.current = false;

        // If closed because of auth error (1008), we might retry with refresh
        if (event.code === 1008) {
          logger.error("WebSocket authentication failed.");
        }

        reconnectAttempts.current += 1;

        if (reconnectTimeout.current) clearTimeout(reconnectTimeout.current);

        // Exponential backoff: 1s, 2s, 4s, 8s, 16s... capped at 30s
        const delay = Math.min(
          1000 * Math.pow(2, reconnectAttempts.current - 1),
          30000,
        );
        logger.debug(`Reconnecting in ${delay}ms...`);

        // Use the ref to call the function recursively
        if (connectRef.current) {
          reconnectTimeout.current = setTimeout(connectRef.current, delay);
        }
      };

      socket.onerror = (error) => {
        logger.error("WebSocket Error:", error);
        // onError usually followed by onClose, so we handle reconnect there
        isConnecting.current = false;
      };
    } catch (error) {
      logger.error("WebSocket Connection Initialization Failed:", error);
      isConnecting.current = false;
      reconnectAttempts.current += 1;
      const delay = Math.min(
        1000 * Math.pow(2, reconnectAttempts.current - 1),
        30000,
      );
      if (connectRef.current) {
        reconnectTimeout.current = setTimeout(connectRef.current, delay);
      }
    }
  }, [user]);

  // Update the ref whenever connect changes
  useEffect(() => {
    connectRef.current = connect;
  }, [connect]);

  useEffect(() => {
    if (user) {
      connect();
    } else {
      // If user logs out, close connection
      if (ws.current) {
        logger.debug("User logged out, closing WebSocket.");
        // Prevent reconnect loop
        ws.current.onopen = null;
        ws.current.onmessage = null;
        ws.current.onerror = null;
        ws.current.onclose = null;
        ws.current.close();
        ws.current = null;
      }
      isConnecting.current = false;
      // Avoid calling setState synchronously
      setTimeout(() => {
        setIsConnected(false);
        setIsFallback(false);
      }, 0);
      reconnectAttempts.current = 0;
      if (reconnectTimeout.current) {
        clearTimeout(reconnectTimeout.current);
      }
    }

    return () => {
      if (ws.current) {
        ws.current.onopen = null;
        ws.current.onmessage = null;
        ws.current.onerror = null;
        ws.current.onclose = null;
        ws.current.close();
        ws.current = null;
      }
      isConnecting.current = false;
      if (reconnectTimeout.current) {
        clearTimeout(reconnectTimeout.current);
      }
    };
  }, [connect, user]);

  const sendMessage = useCallback((msg) => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      logger.debug(`[WebSocket Send] Message action: ${msg.action}`);
      ws.current.send(JSON.stringify(msg));
    } else {
      logger.warn("[WebSocket] Not open, message not sent:", msg);
    }
  }, []);

  const subscribe = useCallback((topic) => {
    logger.debug(`[WebSocket Subscribe] Topic: ${topic}`);
    pendingSubscriptions.current.add(topic);
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify({ action: "subscribe", topic }));
    }
  }, []);

  const unsubscribe = useCallback((topic) => {
    logger.debug(`[WebSocket Unsubscribe] Topic: ${topic}`);
    pendingSubscriptions.current.delete(topic);
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify({ action: "unsubscribe", topic }));
    }
  }, []);

  const reconnect = useCallback(() => {
    reconnectAttempts.current = 0; // Reset attempts on manual reconnect
    setIsFallback(false);
    if (reconnectTimeout.current) clearTimeout(reconnectTimeout.current);
    if (ws.current) {
      ws.current.close(); // This will trigger onclose which triggers reconnect
    } else {
      connect();
    }
  }, [connect]);

  return (
    <WebSocketContext.Provider
      value={{
        isConnected,
        isFallback,
        lastMessage,
        sendMessage,
        subscribe,
        unsubscribe,
        reconnect,
      }}
    >
      {children}
    </WebSocketContext.Provider>
  );
};
