'use client';

import React, {
    createContext,
    useContext,
    useEffect,
    useRef,
    useState,
    useCallback,
} from 'react';
import { useRouter } from 'next/navigation';

import { useAuth } from '@/src/context';
import type { Op } from '@/types';

interface WebSocketContextType {
    addEventListener: (op: Op, handler: (data: any) => void) => void;
    removeEventListener: (op: Op, handler: (data: any) => void) => void;
    isConnected: boolean;
}

const WebSocketContext = createContext<WebSocketContextType | null>(null);

export function WebSocketProvider({ children }: { children: React.ReactNode }) {
    const { token } = useAuth();
    const router = useRouter();

    const [isConnected, setIsConnected] = useState<boolean>(false);
    const eventHandlers = useRef<{ [key in Op]?: ((data: any) => void)[] }>({});

    useEffect(() => {
        if (!token) return;

        const ws = new WebSocket(process.env.NEXT_PUBLIC_WS_BASE!);

        ws.onopen = () => {
            setIsConnected(true);
            ws.send(JSON.stringify({ op: 'AUTHENTICATE', data: { token } }));

            const pingInterval = setInterval(() => {
                ws.send(JSON.stringify({ op: 'PING' }));
            }, 5e3);

            ws.onclose = () => {
                setIsConnected(false);
                clearInterval(pingInterval);
            };
        };

        ws.onmessage = (event) => {
            const { op, data } = JSON.parse(event.data);

            const handlers = eventHandlers.current[op as Op];
            if (!handlers) return;

            console.log(op, handlers);

            handlers.forEach((handler) => handler(data));
        };

        return () => ws.close();
    }, [token]);


    const addEventListener = useCallback(
        (op: Op, handler: (data: any) => void) => {
            if (!eventHandlers.current[op]) {
                eventHandlers.current[op] = [];
            }
            eventHandlers.current[op]!.push(handler);
        },
        []
    );

    const removeEventListener = useCallback(
        (op: Op, handler: (data: any) => void) => {
            if (!eventHandlers.current[op]) return;

            eventHandlers.current[op] = eventHandlers.current[op]!.filter(
                (h) => h !== handler
            );

            if (eventHandlers.current[op]!.length === 0) {
                delete eventHandlers.current[op];
            }
        },
        []
    );

    return (
        <WebSocketContext.Provider
            value={{ addEventListener, removeEventListener, isConnected }}
        >
            {children}
        </WebSocketContext.Provider>
    );
}

export const useWebSocket = () => useContext(WebSocketContext);