import { create } from 'zustand';
import type { Uuid } from '~/lib/api/types';
import type { VoiceServerUpdateServerEvent } from '~/lib/websocket/gateway.types'; // Adjust path
import type { SdpAnswerVoicePayload } from '~/lib/websocket/voice.types';
import { useGatewayWebSocketStore } from './gateway-websocket'; // Adjust path
import { useVoiceWebSocketStore } from './voice-websocket'; // Adjust path

export type WebRTCStatus =
    | "IDLE"
    | "REQUESTING_VOICE_SERVER" // Waiting for Gateway WS to provide voice server info
    | "CONNECTING_VOICE_WS"   // Voice WS connection in progress
    | "NEGOTIATING_SDP"       // SDP Offer/Answer exchange in progress
    | "ICE_GATHERING"         // ICE candidates are being gathered
    | "ICE_CONNECTING"        // ICE connection in progress
    | "CONNECTED"             // WebRTC connection established, media flowing
    | "DISCONNECTED"
    | "FAILED";

interface WebRTCState {
    status: WebRTCStatus;
    peerConnection: RTCPeerConnection | null;
    localStream: MediaStream | null;
    remoteStream: MediaStream | null;
    lastError: string | null;
    currentChannelId: Uuid | null; // To track which channel we are in
    _internalUnsubscribeVoiceWs: ReturnType<typeof useVoiceWebSocketStore.subscribe> | undefined;

    // Actions
    joinVoiceChannel: (serverId: Uuid, channelId: Uuid, localStream: MediaStream) => Promise<void>;
    leaveVoiceChannel: () => void;
    _handleVoiceConnectionInfo: (info: VoiceServerUpdateServerEvent) => void; // Internal, called by GatewayWS

    // For ICE candidate handling (if your server supports trickle ICE via Voice WS)
    // sendIceCandidate: (candidate: RTCIceCandidateInit) => void;
    // _handleIceCandidate: (candidate: RTCIceCandidateInit) => void;
}

// Default ICE server configuration (replace with your own if needed)
const defaultIceServers: RTCIceServer[] = [];

let currentVoiceWsUrl: string | null = null;
let currentVoiceToken: string | null = null;

export const useWebRTCStore = create<WebRTCState>((set, get) => ({
    status: "IDLE",
    peerConnection: null,
    localStream: null,
    remoteStream: null,
    lastError: null,
    currentChannelId: null,
    _internalUnsubscribeVoiceWs: undefined,

    joinVoiceChannel: async (serverId, channelId, localStream) => {
        const { status: gatewayStatus, sendVoiceStateUpdate } = useGatewayWebSocketStore.getState();
        const currentWebRTCStatus = get().status;

        if (currentWebRTCStatus !== "IDLE" && currentWebRTCStatus !== "DISCONNECTED" && currentWebRTCStatus !== "FAILED") {
            console.warn(`WebRTC: Cannot join channel. Current status: ${currentWebRTCStatus}`);
            set({ lastError: "WebRTC: Join attempt while already active or in progress." });
            return;
        }
        if (gatewayStatus !== "CONNECTED") {
            console.error("WebRTC: Gateway WebSocket not connected. Cannot send VOICE_STATE_UPDATE.");
            set({ status: "FAILED", lastError: "Gateway WebSocket not connected." });
            return;
        }

        set({
            status: "REQUESTING_VOICE_SERVER",
            localStream, // Store the local stream
            lastError: null,
            currentChannelId: channelId,
        });

        const payload = { serverId, channelId };
        if (!sendVoiceStateUpdate(payload)) {
            console.error("WebRTC: Failed to send VOICE_STATE_UPDATE via Gateway WS.");
            set({ status: "FAILED", lastError: "Failed to send voice state update." });
            // Revert localStream if needed, or leaveVoiceChannel will clean it up
        } else {
            console.log("WebRTC: VOICE_STATE_UPDATE sent. Waiting for VOICE_CONNECTION_INFO...");
        }
    },

    _handleVoiceConnectionInfo: (info: VoiceServerUpdateServerEvent) => {
        if (get().status !== "REQUESTING_VOICE_SERVER") {
            console.warn("WebRTC: Received VOICE_CONNECTION_INFO in unexpected state:", get().status);
            return; // Or handle error
        }

        console.log("WebRTC: Received voice connection info. Initializing PeerConnection and Voice WS.", info);
        // currentVoiceWsUrl = info.voiceServerUrl;
        currentVoiceWsUrl = "ws://localhost:12345/voice/ws";
        currentVoiceToken = info.data.token;

        const pc = new RTCPeerConnection({ iceServers: defaultIceServers });

        // pc.onicecandidate = (event) => {
        //   if (event.candidate) {
        //     console.log("WebRTC: New ICE candidate generated:", event.candidate);
        //     // IMPORTANT: You need a way to send this to the server via Voice WS
        //     // Example: get().sendIceCandidate(event.candidate.toJSON());
        //     // This requires `sendIceCandidate` on useVoiceWebSocketStore and server support
        //      useVoiceWebSocketStore.getState().sendIceCandidate(event.candidate.toJSON()); // Assuming this exists
        //   } else {
        //     console.log("WebRTC: All ICE candidates have been gathered.");
        //   }
        // };

        pc.oniceconnectionstatechange = () => {
            console.log("WebRTC: ICE connection state change:", pc.iceConnectionState);
            switch (pc.iceConnectionState) {
                case "connected":
                case "completed":
                    set({ status: "CONNECTED", lastError: null });
                    break;
                case "disconnected":
                    // Can sometimes recover, or might lead to "failed"
                    // For now, we might treat as a more final state or wait for "failed"
                    set({ status: "DISCONNECTED", lastError: "ICE connection disconnected." });
                    // Consider calling leaveVoiceChannel or attempting reconnection based on your strategy
                    get().leaveVoiceChannel(); // Simple cleanup on disconnect
                    break;
                case "failed":
                    set({ status: "FAILED", lastError: "ICE connection failed." });
                    get().leaveVoiceChannel(); // Cleanup
                    break;
                case "closed":
                    set(state => (state.status !== "IDLE" ? { status: "DISCONNECTED", lastError: state.lastError || "ICE connection closed." } : {}));
                    break;
                case "new":
                case "checking":
                    set(state => (state.status === "NEGOTIATING_SDP" || state.status === "ICE_GATHERING" ? { status: "ICE_CONNECTING" } : {}));
                    break;
            }
        };

        pc.ontrack = (event) => {
            console.log("WebRTC: Received remote track:", event.track, event.streams);
            if (event.streams && event.streams[0]) {
                set({ remoteStream: event.streams[0] });
            } else {
                // Fallback for older browsers if streams[0] is not available
                const newStream = new MediaStream();
                newStream.addTrack(event.track);
                set({ remoteStream: newStream });
            }
        };

        // Add local tracks to the peer connection
        const localStream = get().localStream;
        if (localStream) {
            localStream.getTracks().forEach(track => {
                try {
                    pc.addTrack(track, localStream);
                    console.log("WebRTC: Added local track:", track.kind);
                } catch (e) {
                    console.error("WebRTC: Error adding local track:", e);
                }
            });
        } else {
            console.warn("WebRTC: No local stream available to add tracks from.");
            // You might want to create a default offer without tracks or handle this case
            // For voice, usually you expect a local audio track.
        }

        set({ peerConnection: pc, status: "CONNECTING_VOICE_WS" });

        // Connect to Voice WebSocket
        const { connect: connectVoiceWs, setCallbacks: setVoiceCallbacks, status: voiceWsStatus } = useVoiceWebSocketStore.getState();

        // Define a handler for SDP Answer from Voice WS
        const handleSdpAnswer = (payload: SdpAnswerVoicePayload) => {
            const currentPc = get().peerConnection;
            if (!currentPc || get().status !== "NEGOTIATING_SDP") {
                console.warn("WebRTC: Received SDP Answer in unexpected state or no PC.", get().status);
                return;
            }
            console.log("WebRTC: Received SDP Answer. Setting remote description.");
            currentPc.setRemoteDescription(new RTCSessionDescription(payload.sdp))
                .then(() => {
                    console.log("WebRTC: Remote description set successfully.");
                    // ICE gathering might already be in progress or starting now.
                    // The oniceconnectionstatechange will handle moving to CONNECTED.
                    set({ status: "ICE_GATHERING" }); // Or directly to ICE_CONNECTING if candidates start flowing
                })
                .catch(err => {
                    console.error("WebRTC: Failed to set remote description:", err);
                    set({ status: "FAILED", lastError: "Failed to set remote SDP answer." });
                    get().leaveVoiceChannel();
                });
        };

        // Define a handler for ICE Candidates from Voice WS (if server sends them)
        const handleVoiceWsIceCandidate = (candidate: RTCIceCandidateInit) => {
            const currentPc = get().peerConnection;
            if (!currentPc) {
                console.warn("WebRTC: Received ICE candidate but no peer connection.");
                return;
            }
            console.log("WebRTC: Received ICE candidate from VoiceWS, adding to PC:", candidate);
            currentPc.addIceCandidate(new RTCIceCandidate(candidate)).catch(e => {
                console.error("WebRTC: Error adding received ICE candidate:", e);
            });
        };


        // Set callbacks for the Voice WS store
        setVoiceCallbacks({
            onSdpAnswer: handleSdpAnswer,
            //   onIceCandidate: handleVoiceWsIceCandidate, // Assuming Voice WS supports this
        });

        // Function to create and send offer
        const createAndSendOffer = async () => {
            const currentPc = get().peerConnection;
            if (!currentPc || get().status !== "NEGOTIATING_SDP") { // Check should be before setting to NEGOTIATING_SDP
                console.warn("WebRTC: Cannot create offer, PC not ready or wrong state.");
                return;
            }
            console.log("WebRTC: Creating SDP Offer...");
            try {
                const offer = await currentPc.createOffer({
                    // Offer to receive audio/video based on what you expect
                    // For voice only:
                    offerToReceiveAudio: true,
                    offerToReceiveVideo: false, // Set to true if you expect video
                });
                await currentPc.setLocalDescription(offer);
                console.log("WebRTC: Local description set. Sending offer via Voice WS.");
                useVoiceWebSocketStore.getState().sendSdpOffer(offer as RTCSessionDescriptionInit); // Cast needed as createOffer returns RTCSessionDescriptionInit
            } catch (err) {
                console.error("WebRTC: Failed to create or send SDP offer:", err);
                set({ status: "FAILED", lastError: "Failed to create/send SDP offer." });
                get().leaveVoiceChannel();
            }
        };


        // Subscribe to Voice WS status changes
        // We need to wait for Voice WS to be 'CONNECTED' before sending offer
        const unsubscribeVoiceWs = useVoiceWebSocketStore.subscribe(
            (newVoiceStatus, oldVoiceStatus) => {
                if (newVoiceStatus.status === "CONNECTED" && get().status === "CONNECTING_VOICE_WS") {
                    console.log("WebRTC: Voice WS connected. Proceeding to SDP negotiation.");
                    set({ status: "NEGOTIATING_SDP" });
                    createAndSendOffer(); // Now create and send offer
                } else if (newVoiceStatus.status === "ERROR" || newVoiceStatus.status === "DISCONNECTED") {
                    if (get().status === "CONNECTING_VOICE_WS" || get().status === "NEGOTIATING_SDP" || get().status === "ICE_GATHERING" || get().status === "ICE_CONNECTING") {
                        console.error("WebRTC: Voice WS disconnected or errored during WebRTC setup.", useVoiceWebSocketStore.getState().lastError);
                        set({ status: "FAILED", lastError: `Voice WebSocket error: ${useVoiceWebSocketStore.getState().lastError || 'Disconnected'}` });
                        get().leaveVoiceChannel(); // Cleanup WebRTC part
                        unsubscribeVoiceWs(); // Clean up subscription
                    }
                }
            }
        );
        // Store unsubscribe function for cleanup in leaveVoiceChannel
        set(state => ({ ...state, _internalUnsubscribeVoiceWs: unsubscribeVoiceWs }));


        // Initiate Voice WS connection
        if (currentVoiceWsUrl && currentVoiceToken) {
            console.log(`WebRTC: Connecting to Voice WS: ${currentVoiceWsUrl}`);
            connectVoiceWs(currentVoiceWsUrl, () => Promise.resolve(currentVoiceToken)); // Token is already a string
        } else {
            console.error("WebRTC: Voice WS URL or Token missing.");
            set({ status: "FAILED", lastError: "Voice WS URL or Token missing." });
        }
    },

    leaveVoiceChannel: () => {
        console.log("WebRTC: Leaving voice channel.");
        const { peerConnection, localStream, _internalUnsubscribeVoiceWs } = get();

        if (_internalUnsubscribeVoiceWs) {
            _internalUnsubscribeVoiceWs(); // Unsubscribe from Voice WS status
        }

        if (peerConnection) {
            peerConnection.getSenders().forEach(sender => {
                if (sender.track) {
                    sender.track.stop();
                }
            });
            peerConnection.getReceivers().forEach(receiver => {
                if (receiver.track) {
                    receiver.track.stop();
                }
            });
            peerConnection.close();
        }
        if (localStream) {
            localStream.getTracks().forEach(track => track.stop());
        }

        // Disconnect Voice WebSocket
        const { status: voiceStatus, disconnect: disconnectVoiceWs } = useVoiceWebSocketStore.getState();
        if (voiceStatus !== "IDLE" && voiceStatus !== "DISCONNECTED") {
            disconnectVoiceWs(true);
        }

        // Optionally, inform Gateway server if needed
        // useGatewayWebSocketStore.getState().sendVoiceStateUpdate({ channelId: null, serverId: get().currentServerId }); // Example

        set({
            status: "IDLE", // Or "DISCONNECTED" if you prefer that as the terminal state after a session
            peerConnection: null,
            localStream: null,
            remoteStream: null,
            lastError: null,
            currentChannelId: null,
            _internalUnsubscribeVoiceWs: undefined,
        });
        currentVoiceWsUrl = null;
        currentVoiceToken = null;
    },

    // Placeholder for sending ICE Candidate via Voice WebSocket
    // sendIceCandidate: (candidate: RTCIceCandidateInit) => {
    //   useVoiceWebSocketStore.getState().sendIceCandidate(candidate); // You'll need to implement sendIceCandidate in useVoiceWebSocketStore
    // },
    // Placeholder for handling ICE Candidate from Voice WebSocket
    // _handleIceCandidate: (candidate: RTCIceCandidateInit) => {
    //   const pc = get().peerConnection;
    //   if (pc && candidate) {
    //     pc.addIceCandidate(new RTCIceCandidate(candidate)).catch(e => {
    //       console.error("WebRTC: Error adding received ICE candidate:", e);
    //     });
    //   }
    // },
}));