import { AccountContext } from "./AccountContext";
import { applyAccessToken, axiosInstance } from "../../axiosApi";
import { setTokens } from "../auth/tokenStorage";
import {
  ConnectionStatus,
  DEFAULT_RECONNECT_CONFIG,
  buildHeartbeatPayload,
  computeBackoffDelay,
  ensureFreshAccessToken,
  isSocketOpen,
  subscribeLifecycleEvents,
} from "./websocketReconnect";

const {
  useEffect,
  useCallback,
  createContext,
  useRef,
  useState,
  useContext,
} = require("react");

const WebSocketContext = createContext();

function WebSocketProvider({ children }) {
  const ws = useRef(null);
  const channels = useRef({});
  const { account } = useContext(AccountContext);
  const accountRef = useRef(account);
  const intentionalCloseRef = useRef(false);
  const reconnectAttemptRef = useRef(0);
  const reconnectTimerRef = useRef(null);
  const heartbeatTimerRef = useRef(null);
  const connectGenerationRef = useRef(0);
  const connectingRef = useRef(false);

  const [socket, setSocket] = useState(null);
  const [connectionStatus, setConnectionStatus] = useState(
    ConnectionStatus.DISCONNECTED,
  );
  /** Increments on each successful open after a drop; MessageContext uses for stream recovery. */
  const [reconnectGeneration, setReconnectGeneration] = useState(0);
  const hadOpenConnectionRef = useRef(false);

  accountRef.current = account;

  const isConnected = connectionStatus === ConnectionStatus.CONNECTED;

  const clearReconnectTimer = useCallback(() => {
    if (reconnectTimerRef.current != null) {
      clearTimeout(reconnectTimerRef.current);
      reconnectTimerRef.current = null;
    }
  }, []);

  const clearHeartbeat = useCallback(() => {
    if (heartbeatTimerRef.current != null) {
      clearInterval(heartbeatTimerRef.current);
      heartbeatTimerRef.current = null;
    }
  }, []);

  const startHeartbeat = useCallback(() => {
    clearHeartbeat();
    heartbeatTimerRef.current = setInterval(() => {
      const current = ws.current;
      if (!isSocketOpen(current)) {
        return;
      }
      try {
        current.send(buildHeartbeatPayload(accountRef.current?.email));
      } catch (err) {
        console.log("WebSocket heartbeat failed", err);
      }
    }, DEFAULT_RECONNECT_CONFIG.heartbeatIntervalMs);
  }, [clearHeartbeat]);

  const subscribe = useCallback((channel, callback) => {
    channels.current[channel] = callback;
  }, []);

  const unsubscribe = useCallback((channel) => {
    delete channels.current[channel];
  }, []);

  const sendMessage = useCallback(
    (message, conversation_id, file, fileType, modelName) => {
      const current = ws.current;
      if (!isSocketOpen(current)) {
        console.log("Error sending message. WebSocket is not open");
        return false;
      }

      const trimmedMessage =
        typeof message === "string" ? message.trim() : "";
      if (!trimmedMessage) {
        console.log("Error sending message. Message text is empty");
        return false;
      }

      if (file) {
        const reader = new FileReader();
        reader.onload = () => {
          const base64File = reader.result?.toString().split(",")[1];
          if (base64File && isSocketOpen(ws.current)) {
            ws.current.send(
              JSON.stringify({
                message: trimmedMessage,
                conversation_id,
                email: accountRef.current?.email,
                file: base64File,
                fileType,
                modelName,
              }),
            );
          }
        };
        reader.readAsDataURL(file);
      } else {
        current.send(
          JSON.stringify({
            message: trimmedMessage,
            conversation_id,
            email: accountRef.current?.email,
            file: null,
            fileType: null,
            modelName,
          }),
        );
      }
      return true;
    },
    [],
  );

  const scheduleReconnect = useCallback(
    (connectFn) => {
      if (intentionalCloseRef.current || !accountRef.current) {
        return;
      }
      if (
        reconnectAttemptRef.current >= DEFAULT_RECONNECT_CONFIG.maxAttempts
      ) {
        console.log("WebSocket reconnect attempts exhausted");
        setConnectionStatus(ConnectionStatus.DISCONNECTED);
        connectingRef.current = false;
        return;
      }

      clearReconnectTimer();
      setConnectionStatus(ConnectionStatus.RECONNECTING);
      const attempt = reconnectAttemptRef.current;
      const delay = computeBackoffDelay(attempt);
      reconnectAttemptRef.current = attempt + 1;

      reconnectTimerRef.current = setTimeout(() => {
        reconnectTimerRef.current = null;
        connectFn({ isRetry: true });
      }, delay);
    },
    [clearReconnectTimer],
  );

  const connect = useCallback(
    async ({ isRetry = false } = {}) => {
      if (!accountRef.current) {
        return;
      }
      if (intentionalCloseRef.current) {
        return;
      }
      if (
        connectingRef.current ||
        isSocketOpen(ws.current) ||
        ws.current?.readyState === WebSocket.CONNECTING
      ) {
        return;
      }

      connectingRef.current = true;
      clearReconnectTimer();
      setConnectionStatus(
        isRetry || hadOpenConnectionRef.current
          ? ConnectionStatus.RECONNECTING
          : ConnectionStatus.CONNECTING,
      );

      // Re-auth before open so REST recovery (#22) and any future WS auth stay valid.
      await ensureFreshAccessToken({
        refreshRequest: async (refresh) => {
          const response = await axiosInstance.post("/token/refresh/", {
            refresh,
          });
          await setTokens(
            response.data.access,
            response.data.refresh ?? refresh,
          );
          applyAccessToken(response.data.access);
          return response.data;
        },
      });

      if (!accountRef.current || intentionalCloseRef.current) {
        connectingRef.current = false;
        return;
      }

      const generation = ++connectGenerationRef.current;
      const url = process.env.REACT_APP_BACKEND_WS_API_BASE_URL;
      if (!url) {
        console.log("REACT_APP_BACKEND_WS_API_BASE_URL is not set");
        connectingRef.current = false;
        setConnectionStatus(ConnectionStatus.DISCONNECTED);
        return;
      }

      try {
        if (ws.current) {
          try {
            ws.current.onopen = null;
            ws.current.onclose = null;
            ws.current.onerror = null;
            ws.current.onmessage = null;
            if (
              ws.current.readyState === WebSocket.OPEN ||
              ws.current.readyState === WebSocket.CONNECTING
            ) {
              ws.current.close();
            }
          } catch {
            /* ignore */
          }
        }

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

        socketInstance.onopen = () => {
          if (generation !== connectGenerationRef.current) {
            return;
          }
          connectingRef.current = false;
          reconnectAttemptRef.current = 0;
          setSocket(socketInstance);
          setConnectionStatus(ConnectionStatus.CONNECTED);
          startHeartbeat();
          if (hadOpenConnectionRef.current) {
            setReconnectGeneration((g) => g + 1);
          }
          hadOpenConnectionRef.current = true;
        };

        socketInstance.onmessage = (message) => {
          if (generation !== connectGenerationRef.current) {
            return;
          }
          const data = message.data;
          if (data === "pong") {
            return;
          }
          try {
            const parsed = JSON.parse(data);
            if (parsed && parsed.type === "pong") {
              return;
            }
          } catch {
            /* not JSON — treat as chat payload */
          }

          const entries = Object.entries(channels.current);
          if (entries.length === 0) {
            console.log("Error: no WebSocket channel subscribers");
            return;
          }
          const [chatChannel] = entries[0];
          if (channels.current[chatChannel]) {
            channels.current[chatChannel](data);
          } else {
            console.log("Error");
          }
        };

        socketInstance.onerror = () => {
          // onclose follows; reconnect is scheduled there
        };

        socketInstance.onclose = () => {
          if (generation !== connectGenerationRef.current) {
            return;
          }
          connectingRef.current = false;
          clearHeartbeat();
          setSocket(null);
          if (intentionalCloseRef.current) {
            setConnectionStatus(ConnectionStatus.DISCONNECTED);
            return;
          }
          scheduleReconnect(connect);
        };
      } catch (err) {
        console.log("WebSocket connect failed", err);
        connectingRef.current = false;
        clearHeartbeat();
        setSocket(null);
        scheduleReconnect(connect);
      }
    },
    [clearHeartbeat, clearReconnectTimer, scheduleReconnect, startHeartbeat],
  );

  const reconnectNow = useCallback(() => {
    if (!accountRef.current || intentionalCloseRef.current) {
      return;
    }
    if (isSocketOpen(ws.current)) {
      return;
    }
    reconnectAttemptRef.current = 0;
    clearReconnectTimer();
    connect({ isRetry: hadOpenConnectionRef.current });
  }, [clearReconnectTimer, connect]);

  useEffect(() => {
    intentionalCloseRef.current = false;
    reconnectAttemptRef.current = 0;
    hadOpenConnectionRef.current = false;
    setReconnectGeneration(0);

    if (!account) {
      intentionalCloseRef.current = true;
      clearReconnectTimer();
      clearHeartbeat();
      connectGenerationRef.current += 1;
      if (ws.current) {
        try {
          ws.current.close();
        } catch {
          /* ignore */
        }
        ws.current = null;
      }
      setSocket(null);
      setConnectionStatus(ConnectionStatus.DISCONNECTED);
      return undefined;
    }

    connect({ isRetry: false });

    const unsubscribeLifecycle = subscribeLifecycleEvents({
      onResume: () => {
        if (!isSocketOpen(ws.current)) {
          reconnectNow();
        }
      },
      onOnline: () => {
        if (!isSocketOpen(ws.current)) {
          reconnectNow();
        }
      },
      onOffline: () => {
        clearReconnectTimer();
        setConnectionStatus((status) =>
          status === ConnectionStatus.CONNECTED
            ? ConnectionStatus.DISCONNECTED
            : status,
        );
      },
    });

    return () => {
      intentionalCloseRef.current = true;
      clearReconnectTimer();
      clearHeartbeat();
      unsubscribeLifecycle();
      connectGenerationRef.current += 1;
      if (ws.current) {
        try {
          ws.current.onopen = null;
          ws.current.onclose = null;
          ws.current.onerror = null;
          ws.current.onmessage = null;
          ws.current.close();
        } catch {
          /* ignore */
        }
        ws.current = null;
      }
      connectingRef.current = false;
      setSocket(null);
      setConnectionStatus(ConnectionStatus.DISCONNECTED);
    };
  }, [
    account,
    clearHeartbeat,
    clearReconnectTimer,
    connect,
    reconnectNow,
  ]);

  return (
    <WebSocketContext.Provider
      value={[
        subscribe,
        unsubscribe,
        socket,
        sendMessage,
        isConnected,
        connectionStatus,
        reconnectGeneration,
        reconnectNow,
      ]}
    >
      {children}
    </WebSocketContext.Provider>
  );
}

export {
  WebSocketContext,
  WebSocketProvider,
  ConnectionStatus,
};
