import { create } from 'zustand';
import type {
    ClientMessage,
    EventServerPayload,
    ServerMessage,
    VoiceStateUpdateClientPayload,
    WebSocketStatus
} from '~/lib/websocket/gateway.types'; // Adjust path as needed
import { useServerListStore } from './server-list'; // Adjust path as needed

// --- Configuration ---
const WEBSOCKET_URL = "ws://localhost:12345/gateway/ws";
const INITIAL_RECONNECT_DELAY = 1000;
const MAX_RECONNECT_DELAY = 30000;
const BACKOFF_FACTOR = 2;
const MAX_RECONNECT_ATTEMPTS = 10;

let isConnectingInProgress = false;
let currentConnectCallId = 0;

function handleBusinessEvent(event: EventServerPayload) {
    console.debug("WS: Received business event:", event.type, event.data);
    switch (event.type) {
        case "ADD_SERVER":
            useServerListStore.getState().addServer(event.data.server);
            break;
        case "REMOVE_SERVER":
            useServerListStore.getState().removeServer(event.data.serverId);
            break;
        // Add other event types from your application
        case "VOICE_SERVER_UPDATE":
            break;
        default:
            // This ensures all event types are handled if EventServerPayload is a strict union
            const _exhaustiveCheck: never = event;
            console.warn("WS: Received unhandled business event type:", _exhaustiveCheck);
            return _exhaustiveCheck;
    }
}

interface GatewayWebSocketState {
    status: WebSocketStatus;
    userId: string | null;
    sessionKey: string | null;
    lastError: string | null;
    connect: (getToken: () => string | null | Promise<string | null>) => void;
    disconnect: (intentional?: boolean) => void;
    sendVoiceStateUpdate: (payload: VoiceStateUpdateClientPayload) => boolean;
}

let socket: WebSocket | null = null;
let pingIntervalId: ReturnType<typeof setInterval> | null = null;
let reconnectTimeoutId: ReturnType<typeof setTimeout> | null = null;
let reconnectAttempts = 0;
let getTokenFunc: (() => string | null | Promise<string | null>) | null = null;
let isIntentionalDisconnect = false;
let serverReportedError = false; // Flag: true if server sent an ERROR message

export const useGatewayWebSocketStore = create<GatewayWebSocketState>((set, get) => {
    const clearPingInterval = () => {
        if (pingIntervalId) clearInterval(pingIntervalId);
        pingIntervalId = null;
    };

    const clearReconnectTimeout = () => {
        if (reconnectTimeoutId) clearTimeout(reconnectTimeoutId);
        reconnectTimeoutId = null;
    };

    const sendWebSocketMessage = (message: ClientMessage): boolean => {
        if (socket?.readyState === WebSocket.OPEN) {
            console.debug("WS: Sending message:", message);
            socket.send(JSON.stringify(message));
            return true;
        }
        console.warn(`WS: Cannot send ${message.type}. WebSocket not open (state: ${socket?.readyState}).`);
        return false;
    };

    const resetConnectingFlag = () => {
        if (isConnectingInProgress) {
            console.debug("WS: Resetting isConnectingInProgress flag.");
            isConnectingInProgress = false;
        }
    };

    const handleOpen = async () => {
        console.log("WS: Connection established. Authenticating...");
        set({ status: "AUTHENTICATING", lastError: null });

        if (!getTokenFunc) {
            console.error("WS: Auth failed. Token getter missing.");
            get().disconnect(false); // Non-intentional, treat as error
            set({ lastError: "Token provider missing for authentication." });
            return;
        }

        try {
            const token = await getTokenFunc();
            if (!token) {
                console.error("WS: Auth failed. No token from getter.");
                isIntentionalDisconnect = true; // Prevent auto-reconnect for this specific failure
                get().disconnect(true);
                set({ lastError: "Authentication token not available." });
                return;
            }
            sendWebSocketMessage({ type: "AUTHENTICATE", data: { token } });
        } catch (error) {
            console.error("WS: Error getting token for auth:", error);
            isIntentionalDisconnect = true;
            get().disconnect(true);
            set({ lastError: "Failed to retrieve authentication token." });
        }
    };

    const handleMessage = async (event: MessageEvent) => {
        try {
            const message = JSON.parse(event.data as string) as ServerMessage;
            console.debug("WS: Received message:", message);

            switch (message.type) {
                case "AUTHENTICATE_ACCEPTED":
                    resetConnectingFlag();
                    reconnectAttempts = 0;
                    clearReconnectTimeout();
                    set({ status: "CONNECTED", userId: message.data.userId, sessionKey: message.data.sessionKey, lastError: null });
                    console.log(`WS: Authenticated as user ${message.data.userId}.`);

                    break;
                case "AUTHENTICATE_DENIED":
                    resetConnectingFlag();
                    console.warn("WS: Authentication denied by server.");
                    isIntentionalDisconnect = true; // Prevent auto-reconnect
                    serverReportedError = false; // This is an auth denial, not a runtime server error
                    get().disconnect(true); // Disconnect, don't retry
                    set({ status: "ERROR", lastError: "Authentication denied by server." });
                    break;
                case "EVENT":
                    handleBusinessEvent(message.data.event);
                    break;
                case "ERROR":
                    resetConnectingFlag();
                    console.error(`WS: Server reported error. Code: ${message.data.code}`);
                    serverReportedError = true; // CRITICAL: Set flag
                    isIntentionalDisconnect = true; // Treat as a definitive stop from server
                    set({
                        status: "ERROR",
                        lastError: `Server error (${message.data.code})`,
                        userId: null, // Clear user session on server error
                        sessionKey: null, // Clear session key on server error
                    });
                    // Server is expected to close the connection. `onclose` will use `serverReportedError`.
                    // Ping interval will be cleared by `cleanupConnection` via `onclose`.
                    break;
                default:
                    const _exhaustiveCheck: never = message;
                    console.warn("WS: Received unknown server message type:", _exhaustiveCheck);
                    set({ lastError: "Received unknown message type from server." });
            }
        } catch (error) {
            console.error("WS: Failed to parse or handle server message:", error, "Raw data:", event.data);
            set({ lastError: "Failed to parse server message." });
            // This is a client-side parsing error. We might want to disconnect if this happens.
            // serverReportedError = true; // Consider this a fatal client error
            // isIntentionalDisconnect = true;
            // get().disconnect(false); // Or true depending on desired reconnect behavior
        }
    };

    const handleError = (event: Event) => {
        resetConnectingFlag();
        // This event is often vague and usually followed by onclose.
        console.error("WS: WebSocket error event occurred:", event);
        // Avoid changing status directly if onclose will handle it, to prevent state flickering.
        // But, set lastError as it might be the only indication of a problem if onclose is not detailed.
        set(state => ({ lastError: state.lastError || `WebSocket error: ${event.type || 'Unknown error'}` }));
        // `onclose` will call `cleanupConnection`.
    };

    const handleClose = (event: CloseEvent) => {
        const wasInProgress = isConnectingInProgress;
        resetConnectingFlag();
        console.log(
            `WS: Connection closed. Code: ${event.code}, Reason: "${event.reason || 'N/A'}", Clean: ${event.wasClean}, Intentional: ${isIntentionalDisconnect}, ServerError: ${serverReportedError}, WasInProgress: ${wasInProgress}`
        );

        const wasNormalClosure = event.code === 1000 || event.code === 1001; // 1001: Going Away
        const shouldAttemptReconnect = !isIntentionalDisconnect && !serverReportedError && !wasNormalClosure;

        const previousStatus = get().status;
        cleanupConnection(shouldAttemptReconnect);

        if (!shouldAttemptReconnect) {
            console.log("WS: No reconnect attempt scheduled.");
            if (serverReportedError) {
                // State already set to ERROR by handleMessage or will be set by cleanup.
                if (get().status !== 'ERROR') set({ status: "ERROR", userId: null, sessionKey: null }); // Ensure it
            } else if (isIntentionalDisconnect && previousStatus !== "ERROR") {
                set({ status: "DISCONNECTED", userId: null, sessionKey: null });
            } else if (wasNormalClosure && previousStatus !== "ERROR") {
                set({ status: "DISCONNECTED", userId: null, sessionKey: null });
            } else if (previousStatus !== "ERROR" && previousStatus !== "DISCONNECTED" && previousStatus !== "IDLE") {
                set({ status: "DISCONNECTED", userId: null, sessionKey: null, lastError: get().lastError || "Connection closed unexpectedly." });
            }
        }
        // If shouldAttemptReconnect, cleanupConnection -> scheduleReconnect sets status to RECONNECTING.
    };

    const cleanupConnection = (attemptReconnect: boolean) => {
        resetConnectingFlag();
        console.debug(`WS: Cleaning up connection. Attempt reconnect: ${attemptReconnect}`);
        clearPingInterval();
        clearReconnectTimeout();

        if (socket) {
            socket.onopen = null; socket.onmessage = null; socket.onerror = null; socket.onclose = null;
            if (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING) {
                try { socket.close(1000, "Client cleanup"); }
                catch (e) { console.warn("WS: Error closing socket during cleanup:", e); }
            }
            socket = null;
        }

        if (attemptReconnect) {
            if (!serverReportedError) { // Don't retry if server explicitly said ERROR
                scheduleReconnect();
            } else {
                console.warn("WS: Reconnect attempt skipped due to server-reported error.");
                if (get().status !== 'ERROR') {
                    set({ status: "ERROR", userId: null, sessionKey: null, lastError: get().lastError || "Server error, stopping." });
                }
            }
        }
    };

    const scheduleReconnect = () => {
        if (get().status === "ERROR" && serverReportedError) { // Double check
            console.log("WS: Reconnect inhibited by server-reported error state.");
            return;
        }
        if (reconnectAttempts >= MAX_RECONNECT_ATTEMPTS) {
            console.warn(`WS: Max reconnect attempts (${MAX_RECONNECT_ATTEMPTS}) reached.`);
            set({ status: "ERROR", lastError: "Reconnection failed after multiple attempts.", userId: null, sessionKey: null });
            return;
        }

        const delay = Math.min(INITIAL_RECONNECT_DELAY * Math.pow(BACKOFF_FACTOR, reconnectAttempts), MAX_RECONNECT_DELAY);
        reconnectAttempts++;
        console.log(`WS: Scheduling reconnect attempt ${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} in ${delay / 1000}s...`);
        set(state => ({ status: "RECONNECTING", lastError: state.lastError }));

        reconnectTimeoutId = setTimeout(() => {
            if (get().status !== "RECONNECTING") { // Status might have changed (e.g. explicit disconnect)
                console.log("WS: Reconnect attempt cancelled (status changed).");
                return;
            }
            _connect();
        }, delay);
    };

    const _connect = async () => { // Internal connect, used for retries and initial
        const callId = currentConnectCallId;
        if (socket && (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING)) {
            console.warn(`WS (_connect [${callId}]): Called but WebSocket already open/connecting.`);
            if (socket.readyState === WebSocket.OPEN && get().status === 'CONNECTED') resetConnectingFlag();
            return;
        }

        if (isConnectingInProgress) {
            console.warn(`WS (_connect [${callId}]): Aborted, another connection attempt is already in progress.`);
            return;
        }

        // Critical: If serverReportedError is true, _connect (retry mechanism) should not proceed.
        // Only an explicit call to the public `connect()` (which resets serverReportedError) should bypass this.
        if (serverReportedError) {
            resetConnectingFlag();
            console.warn("WS: _connect aborted due to server-reported error. Explicit connect() needed to retry.");
            if (get().status !== 'ERROR') set({ status: 'ERROR', lastError: get().lastError || 'Server reported an error.' });
            return;
        }

        clearReconnectTimeout();

        if (!getTokenFunc) {
            console.error("WS: Cannot connect. Token getter missing.");
            set({ status: "ERROR", lastError: "Token provider missing." }); // Terminal for this attempt cycle
            return;
        }

        console.log(`WS (_connect [${callId}]): Attempting to connect to ${WEBSOCKET_URL}... (Overall Attempt ${reconnectAttempts + 1})`);
        isConnectingInProgress = true;

        try { // Pre-flight check for token availability
            const token = await getTokenFunc();
            if (!token) {
                console.warn("WS: No token available. Aborting connection attempt cycle.");
                set({ status: "DISCONNECTED", lastError: "Authentication token unavailable.", userId: null, sessionKey: null });
                isIntentionalDisconnect = true; // Stop retrying if token is consistently unavailable
                return;
            }
        } catch (error) {
            console.error("WS: Error getting token before connecting:", error);
            set({ status: "ERROR", lastError: "Failed to retrieve authentication token.", userId: null, sessionKey: null });
            isIntentionalDisconnect = true; // Stop retrying for this
            return;
        }

        console.log(`WS: Attempting to connect to ${WEBSOCKET_URL}... (Attempt ${reconnectAttempts + 1})`);
        // Preserve lastError if we are RECONNECTING, otherwise clear it for a fresh CONNECTING attempt
        set(state => ({ status: "CONNECTING", lastError: state.status === 'RECONNECTING' ? state.lastError : null }));

        try {
            socket = new WebSocket(WEBSOCKET_URL);
            // Flags (isIntentionalDisconnect, serverReportedError) are reset by public `connect()`
            socket.onopen = handleOpen;
            socket.onmessage = handleMessage;
            socket.onerror = handleError;
            socket.onclose = handleClose;
        } catch (error) {
            console.error("WS: Failed to create WebSocket instance:", error);
            set({ status: "ERROR", lastError: "Failed to initialize WebSocket connection." });
            socket = null;
            resetConnectingFlag();
            cleanupConnection(!isIntentionalDisconnect && !serverReportedError); // Attempt reconnect if appropriate
        }
    };

    return {
        status: "IDLE",
        userId: null,
        sessionKey: null,
        lastError: null,

        connect: (newGetTokenFunc) => {
            const localCallId = ++currentConnectCallId;
            console.log(`WS: Explicit connect requested (Call ID: ${localCallId}). Current status: ${get().status}, InProgress: ${isConnectingInProgress}`);

            // If already connected and socket is open, do nothing.
            if (get().status === "CONNECTED" && socket?.readyState === WebSocket.OPEN) {
                console.warn(`WS (connect [${localCallId}]): Already connected.`);
                resetConnectingFlag(); // Ensure it's false if we are truly connected.
                return;
            }

            // If a connection is actively being established by another call, log and return.
            // Allow if status is IDLE/DISCONNECTED/ERROR, as these are states where a new attempt is valid.
            if (isConnectingInProgress && !['IDLE', 'DISCONNECTED', 'ERROR', 'RECONNECTING'].includes(get().status)) {
                console.warn(`WS (connect [${localCallId}]): Connect called while a connection attempt is already in progress and status is ${get().status}.`);
                return;
            }
            // If status is CONNECTING/AUTHENTICATING but isConnectingInProgress is false (e.g., after a bug),
            // allow cleanup and new attempt.

            // Reset flags for a *new* explicit connection sequence
            isIntentionalDisconnect = false;
            serverReportedError = false;
            reconnectAttempts = 0;
            // isConnectingInProgress will be set by _connect if it proceeds.
            // Resetting it here might be too early if _connect has checks that depend on previous state of isConnectingInProgress.
            // Instead, ensure _connect and its outcomes (open, error, close) reliably reset it.

            clearReconnectTimeout();

            if (socket) { // If there's an old socket (e.g. from a failed/interrupted attempt)
                console.warn(`WS (connect [${localCallId}]): Explicit connect called with existing socket (state: ${socket.readyState}). Cleaning up old one.`);
                cleanupConnection(false); // This will call resetConnectingFlag()
            }
            // At this point, isConnectingInProgress should be false due to cleanupConnection or initial state.

            getTokenFunc = newGetTokenFunc;
            _connect(); // _connect will set isConnectingInProgress = true if it starts.
        },

        disconnect: (intentional = true) => {
            const callId = currentConnectCallId; // For logging
            console.log(`WS (disconnect [${callId}]): Disconnect requested. Intentional: ${intentional}, InProgress: ${isConnectingInProgress}`);
            isIntentionalDisconnect = intentional;
            if (intentional) {
                serverReportedError = false;
            }
            // No matter what, a disconnect means any "connecting in progress" is now void.
            resetConnectingFlag(); // Crucial: stop any perceived ongoing connection attempt.
            reconnectAttempts = 0;

            const currentLastError = get().lastError;
            const wasServerError = serverReportedError || (currentLastError && currentLastError.startsWith("Server error"));
            const wasAuthDenied = currentLastError === "Authentication denied by server.";

            cleanupConnection(false); // Clean up, do not attempt reconnect

            set({
                status: "DISCONNECTED",
                userId: null,
                sessionKey: null,
                // Preserve significant errors (server error, auth denied) even on user disconnect.
                // Otherwise, set a generic disconnect message.
                lastError: (wasServerError || wasAuthDenied) ? currentLastError : (intentional ? "User disconnected" : "Disconnected"),
            });
        },
        sendVoiceStateUpdate: (payload: VoiceStateUpdateClientPayload) => {
            return sendWebSocketMessage({ type: "VOICE_STATE_UPDATE", data: payload });
        }
    };
});

// Optional: Online/Offline listeners. Consider placing in a React component
// that can supply `getTokenFunc` more naturally upon reconnection.
if (typeof window !== 'undefined') {
    const tryAutoConnectOnline = () => {
        const { status, connect } = useGatewayWebSocketStore.getState();
        console.log("WS: Browser online event detected.");
        if (status === 'DISCONNECTED' || status === 'ERROR' || status === 'RECONNECTING') {
            if (!serverReportedError) { // Don't auto-connect if last state was a server-reported fatal error
                console.log("WS: Attempting to connect after coming online.");
                if (getTokenFunc) { // Check if we still have a token getter
                    connect(getTokenFunc);
                } else {
                    console.warn("WS: Cannot auto-connect on 'online': token getter not available. App needs to call connect() again.");
                }
            } else {
                console.log("WS: Browser online, but previous server-reported error prevents auto-reconnect. Manual connect needed.");
            }
        }
    };
    window.addEventListener('online', tryAutoConnectOnline);

    const handleAppClose = () => { // Best-effort for page unload
        const { disconnect, status } = useGatewayWebSocketStore.getState();
        if (status !== "IDLE" && status !== "DISCONNECTED") {
            disconnect(true);
        }
    };
    window.addEventListener('beforeunload', handleAppClose);
}