import { useState, useEffect, useRef, useCallback } from 'react';

export function useWebSocket(sessionId: string | null) {
    const [isConnected, setIsConnected] = useState(false);
    const [isSynchronized, setIsSynchronized] = useState(false);
    const [lastMessage, setLastMessage] = useState<any | null>(null);
    const [clientId, setClientId] = useState<string | null>(null);
    const ws = useRef<WebSocket | null>(null);

    const connect = useCallback(() => {
        if (!sessionId || ws.current) return;

        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${window.location.host}/ag-beats?sessionId=${sessionId}`;
        const socket = new WebSocket(wsUrl);

        socket.onopen = () => {
            console.log('WebSocket connected');
            setIsConnected(true);
            // Request the full session state upon connecting
            socket.send(JSON.stringify({ type: 'get_state' }));
        };

        socket.onmessage = (event) => {
            const message = JSON.parse(event.data);
            if (message.type === 'welcome') {
                setClientId(message.payload.clientId);
            } else if (message.type === 'session_state') {
                // This is the full state, apply it and mark as synchronized
                setLastMessage({ type: 'state', payload: message.payload });
                setIsSynchronized(true);
            } else {
                setLastMessage(message);
            }
        };

        socket.onclose = () => {
            console.log('WebSocket disconnected');
            setIsConnected(false);
            setIsSynchronized(false);
            ws.current = null;
        };

        socket.onerror = (error) => {
            console.error('WebSocket error:', error);
            socket.close();
        };

        ws.current = socket;
    }, [sessionId]);

    useEffect(() => {
        connect();
        return () => {
            ws.current?.close();
        };
    }, [connect]);

    const sendMessage = (message: any) => {
        if (ws.current && ws.current.readyState === WebSocket.OPEN) {
            ws.current.send(JSON.stringify(message));
        }
    };

    return { isConnected, isSynchronized, lastMessage, sendMessage, clientId };
}