import React, { useEffect, useRef } from "react";
import { useBaseStore } from "../stores/base";
import { useChatStore } from "../stores/chats";
import { useWsStore } from "../stores/ws";
import { Opcode, ReadyOp } from "../types/opcodes";
import { WsMessage } from "../types/types";
import { Message } from "../types/discord";
import { URL } from "react-native-url-polyfill";

export const WebSocketProvider = ({ children }: { children: React.ReactNode }) => {
  const token = useBaseStore((s) => s.token);
  const baseUrl = useBaseStore((s) => s.baseUrl);
  const setUser = useBaseStore((s) => s.setUser);
  const clearAuth = useBaseStore((s) => s.clearAuth);

  const setConnected = useWsStore((s) => s.setConnected);
  const setPrivateChannels = useChatStore((s) => s.setPrivateChannels);
  const addMessage = useChatStore((s) => s.addMessage);

  const tokenRef = useRef(token);
  const baseUrlRef = useRef(baseUrl);
  
  const wsRef = useRef<WebSocket | null>(null);
  const heartbeatRef = useRef<ReturnType<typeof setInterval> | null>(null);
  const retryTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const shouldConnectRef = useRef(false);

  useEffect(() => {
    tokenRef.current = token;
    baseUrlRef.current = baseUrl;
  }, [token, baseUrl]);

  const cleanup = () => {
    shouldConnectRef.current = false;

    if (retryTimeoutRef.current) {
      clearTimeout(retryTimeoutRef.current);
      retryTimeoutRef.current = null;
    }

    if (wsRef.current) {
      try {
        stopHeartbeat();
        wsRef.current.onopen = null;
        wsRef.current.onmessage = null;
        wsRef.current.onerror = null;
        wsRef.current.onclose = null;
        wsRef.current.close();
      } catch (e) {
        // ignore
      }
      wsRef.current = null;
    }

    setConnected(false);
  };

  const getWebSocketUrlFrom = (base: string) => {
    const url = new URL(base);
    url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
    url.pathname = "/ws";
    return url.toString();
  };

  const startHeartbeat = (interval: number) => {
    if (heartbeatRef.current) clearInterval(heartbeatRef.current);
    heartbeatRef.current = setInterval(() => {
      if (wsRef.current?.readyState === WebSocket.OPEN) {
        wsRef.current.send(JSON.stringify({ op: Opcode.HEARTBEAT, d: Date.now() }));
      }
    }, interval);
  };

  const stopHeartbeat = () => {
    if (heartbeatRef.current) {
      clearInterval(heartbeatRef.current);
      heartbeatRef.current = null;
    }
  };

  const handleDispatch = (eventType: string, payload: any) => {
    switch (eventType) {
      case "READY": {
        const readyPayload = payload as ReadyOp;
        if (!readyPayload) return;

        setUser(readyPayload.user);
        setPrivateChannels(readyPayload.private_channels);
        break;
      }
      case "MESSAGE_CREATE": {
        const msg = payload as Message;
        addMessage(msg.channel_id, msg);
        break;
      }
      default:
        console.debug("Unhandled dispatch:", eventType, payload);
    }
  };

  const handleMessage = (raw: string) => {
    let parsed: WsMessage;
    try {
      parsed = JSON.parse(raw);
    } catch {
      console.warn("Invalid WS message:", raw);
      return;
    }

    const { op, d, t } = parsed;

    switch (op) {
      case Opcode.HELLO:
        const interval = d?.heartbeat_interval ?? 30000;
        startHeartbeat(interval);
        break;
      case Opcode.DISPATCH:
        if (t) handleDispatch(t, d);
        break;
      case Opcode.INVALID_SESSION:
        clearAuth();
        useChatStore.getState().resetStore()
        break;
      default:
        console.debug("Unhandled op:", op, d);
    }
  };

  const connect = (token: string, baseUrl: string) => {
    let url: string;
    try {
      url = getWebSocketUrlFrom(baseUrl);
    } catch (err) {
      console.warn("Invalid baseUrl", err);
      return;
    }

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

    ws.onopen = () => {
      if (!shouldConnectRef.current) {
        ws.close();
        return;
      }
      setConnected(true);
      try {
        ws.send(JSON.stringify({ op: Opcode.IDENTIFY, d: { token } }));
      } catch (err) {
        console.warn("Failed to send IDENTIFY", err);
      }
    };

    ws.onmessage = (ev) => {
      if (!shouldConnectRef.current) return;
      const data = typeof ev.data === "string" ? ev.data : new TextDecoder().decode(ev.data as any);
      handleMessage(data);
    };

    const handleDisconnect = () => {
      stopHeartbeat();
      useChatStore.getState().resetStore();
      setConnected(false);

      if (!shouldConnectRef.current) return;
      retryTimeoutRef.current = setTimeout(() => {
        connect(tokenRef.current!, baseUrlRef.current!);
      }, 1000 + Math.random() * 1000);
    };

    ws.onerror = handleDisconnect;
    ws.onclose = handleDisconnect;
  };

  useEffect(() => {
    cleanup();
    if (!token || !baseUrl) return;

    shouldConnectRef.current = true;
    connect(token, baseUrl);

    return () => {
      shouldConnectRef.current = false;
      cleanup();
    };
  }, [token, baseUrl]);

  return <>{children}</>;
};