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 setConnected = useWsStore((s) => s.setConnected);
  const setUser = useBaseStore((s) => s.setUser);
  const setPrivateChannels = useChatStore((s) =>  s.setPrivateChannels)
  const addMessage = useChatStore((s) =>  s.addMessage)

  const token = useBaseStore((s) => s.token);
  const baseUrl = useBaseStore((s) => s.baseUrl);
  
  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 {
        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;
    }

    if (heartbeatRef.current) {
      clearInterval(heartbeatRef.current);
      heartbeatRef.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 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);
      }
      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;
        if (heartbeatRef.current) clearInterval(heartbeatRef.current);

        heartbeatRef.current = setInterval(() => {
          if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;
          wsRef.current.send(JSON.stringify({ op: Opcode.HEARTBEAT, d: Date.now() }));
        }, interval);
        break;
      }
      case Opcode.DISPATCH:
        if (t) handleDispatch(t, d);
        break;
      case Opcode.RECONNECT:
        wsRef.current?.close();
        break;
      case Opcode.INVALID_SESSION:
        useBaseStore.getState().clearAuth();
        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, skipping WebSocket connect", err);
      return;
    }

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

    ws.onopen = () => {
      if (!shouldConnectRef.current) {
        ws.close();
        return;
      }
      setConnected(true);
      ws.send(JSON.stringify({ op: Opcode.IDENTIFY, d: { token } }));
    };

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

    ws.onerror = (ev) => {
      console.warn("WebSocket error", ev);
      setConnected(false);
    };

    ws.onclose = () => {
      setConnected(false);
      if (heartbeatRef.current) {
        clearInterval(heartbeatRef.current);
        heartbeatRef.current = null;
      }
      
      if (!shouldConnectRef.current) return;
      retryTimeoutRef.current = setTimeout(() => {
        connect(tokenRef.current!, baseUrlRef.current!);
      }, 1000 + Math.random() * 1000);
    };
  };

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

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

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

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