import { create } from 'zustand';
import type { SdpAnswerVoicePayload, VoiceClientMessage, VoiceServerMessage, WebSocketStatus } from '~/lib/websocket/voice.types';

// Callbacks for application-specific events
interface VoiceWebSocketCallbacks {
  onSdpAnswer?: (payload: SdpAnswerVoicePayload) => void;
  // onBrowserOnline?: () => void; // Optional: If app wants to be notified by 'online' event
}

// --- Store State ---
export interface VoiceWebSocketState {
  status: WebSocketStatus;
  lastError: string | null;
  connect: (
    webSocketUrl: string,
    getToken: () => string | null | Promise<string | null>
  ) => void;
  disconnect: (intentional?: boolean) => void;
  sendSdpOffer: (sdp: RTCSessionDescriptionInit) => boolean;
  setCallbacks: (callbacks: VoiceWebSocketCallbacks) => void;
}

// --- Module-level state (managed by Zustand store closure) ---
let socket: WebSocket | null = null;
let getTokenFunc: (() => string | null | Promise<string | null>) | null = null;
let currentWebSocketUrl: string | null = null;

let isIntentionalDisconnect = false;
let serverReportedError = false; // Flag: true if server sent an ERROR message
let isConnectingInProgress = false; // Flag to prevent multiple concurrent connect attempts
let currentConnectCallId = 0; // For debugging concurrent calls

let voiceCallbacks: VoiceWebSocketCallbacks = {};

export const useVoiceWebSocketStore = create<VoiceWebSocketState>((set, get) => {

  const sendWebSocketMessage = (message: VoiceClientMessage): boolean => {
    if (socket?.readyState === WebSocket.OPEN) {
      console.debug("VoiceWS: Sending message:", message.type, message.data);
      try {
        socket.send(JSON.stringify(message));
        return true;
      } catch (error) {
        console.error("VoiceWS: Error sending message:", error, message);
        // This is a critical error during send, transition to ERROR
        set({ status: "ERROR", lastError: "Failed to send message due to a WebSocket error." });
        cleanupConnectionInternals(); // Clean up the broken socket
        return false;
      }
    }
    console.warn(`VoiceWS: Cannot send ${message.type}. WebSocket not open (state: ${socket?.readyState}).`);
    return false;
  };

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

  // Cleans up socket event handlers and closes the socket if open/connecting.
  // This function does NOT change the Zustand state directly, relying on callers or event handlers (like onclose) to do that.
  const cleanupConnectionInternals = () => {
    console.debug(`VoiceWS: Cleaning up connection internals.`);
    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("VoiceWS: Error closing socket during cleanup:", e); }
      }
      socket = null;
    }
  };

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

    if (!getTokenFunc) {
      console.error("VoiceWS: Auth failed. Token getter missing.");
      resetConnectingFlag();
      set({ status: "ERROR", lastError: "Token provider missing for authentication." });
      cleanupConnectionInternals();
      return;
    }

    try {
      const token = await getTokenFunc();
      if (!token) {
        console.error("VoiceWS: Auth failed. No token from getter.");
        isIntentionalDisconnect = true; // So handleClose doesn't treat as "unexpected"
        resetConnectingFlag();
        set({ status: "ERROR", lastError: "Authentication token not available." });
        cleanupConnectionInternals();
        return;
      }
      // If sendWebSocketMessage fails here, it will set state to ERROR and cleanup.
      if (!sendWebSocketMessage({ type: "AUTHENTICATE", data: { token } })) {
        // Error state already set by sendWebSocketMessage if it failed critically
        resetConnectingFlag(); // Ensure flag is reset
      }
    } catch (error) {
      console.error("VoiceWS: Error getting token for auth:", error);
      isIntentionalDisconnect = true;
      resetConnectingFlag();
      set({ status: "ERROR", lastError: "Failed to retrieve authentication token." });
      cleanupConnectionInternals();
    }
  };

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

      switch (message.type) {
        case "AUTHENTICATE_ACCEPTED":
          resetConnectingFlag();
          set({ status: "CONNECTED", lastError: null });
          console.log(`VoiceWS: Authenticated successfully.`);
          break;
        case "AUTHENTICATE_DENIED":
          resetConnectingFlag();
          const reason = message.data?.reason || "No reason provided";
          console.warn(`VoiceWS: Authentication denied by server. Reason: ${reason}`);
          isIntentionalDisconnect = true; // Server is denying, so it's "intentional" from that PoV
          serverReportedError = false; // This is an auth failure, not a general server runtime error
          set({ status: "ERROR", lastError: `Authentication denied: ${reason}` });
          // Server should close the connection. We can prompt it.
          if (socket) socket.close(1000, "Authentication Denied");
          // cleanupConnectionInternals() will be called by handleClose
          break;
        case "SDP_ANSWER":
          if (voiceCallbacks.onSdpAnswer) {
            voiceCallbacks.onSdpAnswer(message.data);
          } else {
            console.warn("VoiceWS: Received SDP_ANSWER but no handler is registered.");
          }
          break;
        case "ERROR":
          resetConnectingFlag();
          const errCode = message.data.code;
          const errMsg = message.data.message || "Unknown server error";
          console.error(`VoiceWS: Server reported error. Code: ${errCode}, Message: "${errMsg}"`);
          serverReportedError = true;
          isIntentionalDisconnect = true; // Server error implies server wants to stop.
          set({
            status: "ERROR",
            lastError: `Server error (${errCode}): ${errMsg}`,
          });
          // Server should close the connection. We can prompt it.
          if (socket) socket.close(1000, "Server Reported Error");
          // cleanupConnectionInternals() will be called by handleClose
          break;
        default:
          const _exhaustiveCheck: never = message;
          console.warn("VoiceWS: Received unknown server message type:", _exhaustiveCheck);
          set({ status: "ERROR", lastError: "Received unknown message type from server." });
          if (socket) socket.close(1000, "Unknown message type");
          return _exhaustiveCheck;
      }
    } catch (error) {
      console.error("VoiceWS: Failed to parse or handle server message:", error, "Raw data:", event.data);
      resetConnectingFlag();
      set({ status: "ERROR", lastError: "Failed to parse server message." });
      if (socket) socket.close(1000, "Message parsing error");
      // cleanupConnectionInternals() will be called by handleClose
    }
  };

  const handleError = (event: Event) => {
    console.error("VoiceWS: WebSocket error event occurred:", event);
    // This event often precedes `onclose`.
    // `isConnectingInProgress` should be reset by `onclose`.
    // Set lastError if not already set by a more specific server message or a connection setup failure.
    set(state => ({
      lastError: state.lastError || `WebSocket error: ${event.type || 'Unknown error'}`
    }));
    // Do not change status to ERROR here; onclose will give the definitive closure reason.
    // However, ensure isConnectingInProgress is reset if this is a terminal error before onclose.
    if (socket?.readyState === WebSocket.CLOSING || socket?.readyState === WebSocket.CLOSED) {
      resetConnectingFlag();
    }
  };

  const handleClose = (event: CloseEvent) => {
    const wasInProgress = isConnectingInProgress; // Capture before reset
    resetConnectingFlag(); // Connection process is definitely over.

    console.log(
      `VoiceWS: Connection closed. Code: ${event.code}, Reason: "${event.reason || 'N/A'}", Clean: ${event.wasClean}, Intentional: ${isIntentionalDisconnect}, ServerError: ${serverReportedError}, WasInProgress: ${wasInProgress}`
    );

    cleanupConnectionInternals(); // Ensure all handlers are detached and socket is nulled.

    const currentStatus = get().status;
    const currentError = get().lastError;

    // If status was already set to ERROR by handleMessage (e.g. AUTH_DENIED, server ERROR),
    // or by a failure during connect/auth phases, preserve that specific error.
    if (currentStatus === "ERROR") {
      // If lastError is generic like "Failed to initialize...", update with close event info if more specific.
      if (currentError === "Failed to initialize WebSocket connection." || !currentError) {
        set({ lastError: `Connection closed: Code ${event.code}, Reason: "${event.reason || 'N/A'}"` });
      }
      // Otherwise, the specific error (like auth denial) is already set and should be kept.
      return;
    }

    // If an explicit server error message was received and handled
    if (serverReportedError) {
      set({ status: "ERROR", lastError: currentError || `Server error led to closure (Code: ${event.code})` });
    }
    // If disconnect was called explicitly, or an auth denial occurred (isIntentionalDisconnect=true)
    else if (isIntentionalDisconnect) {
      set({
        status: "DISCONNECTED",
        lastError: currentError || (event.wasClean ? "Disconnected" : `Disconnected: Code ${event.code}, Reason: "${event.reason || 'N/A'}"`)
      });
    }
    // Otherwise, it's an unexpected closure (e.g., network drop, server crash without specific ERROR message)
    else {
      set({
        status: "ERROR",
        lastError: `Connection lost unexpectedly: Code ${event.code}, Reason: "${event.reason || 'N/A'}"`
      });
    }
  };

  const _connect = async () => {
    const callId = currentConnectCallId;

    // This check is crucial: if a server error occurred, only a new explicit `connect()`
    // (which resets serverReportedError via the public connect method) should allow a new attempt.
    if (serverReportedError) {
      console.warn("VoiceWS: _connect aborted due to prior server-reported error. Call public connect() to retry.");
      // State should already be ERROR. Ensure connecting flag is false.
      resetConnectingFlag();
      return;
    }

    if (!currentWebSocketUrl) { // Should be set by public connect()
      console.error("VoiceWS: Cannot connect. WebSocket URL missing.");
      set({ status: "ERROR", lastError: "WebSocket URL not provided." });
      resetConnectingFlag();
      return;
    }
    if (!getTokenFunc) { // Should be set by public connect()
      console.error("VoiceWS: Cannot connect. Token getter missing.");
      set({ status: "ERROR", lastError: "Token provider missing." });
      resetConnectingFlag();
      return;
    }

    console.log(`VoiceWS (_connect [${callId}]): Attempting to connect to ${currentWebSocketUrl}...`);
    isConnectingInProgress = true; // Set flag: connection attempt now in progress.

    try {
      const token = await getTokenFunc(); // Pre-flight check for token
      if (!token) {
        console.warn("VoiceWS: No token available during connection attempt. Aborting.");
        set({ status: "ERROR", lastError: "Authentication token unavailable for connection." });
        isIntentionalDisconnect = true; // To guide handleClose if something unexpected happens next
        resetConnectingFlag();
        return;
      }
    } catch (error) {
      console.error("VoiceWS: Error getting token before connecting:", error);
      set({ status: "ERROR", lastError: "Failed to retrieve authentication token before connecting." });
      isIntentionalDisconnect = true;
      resetConnectingFlag();
      return;
    }

    set({ status: "CONNECTING", lastError: null }); // Fresh connection attempt

    try {
      socket = new WebSocket(currentWebSocketUrl);
      socket.onopen = handleOpen;
      socket.onmessage = handleMessage;
      socket.onerror = handleError;
      socket.onclose = handleClose;
    } catch (error) {
      console.error(`VoiceWS: Failed to create WebSocket instance for ${currentWebSocketUrl}:`, error);
      set({ status: "ERROR", lastError: "Failed to initialize WebSocket connection." });
      socket = null;
      resetConnectingFlag(); // Reset before cleanup if WebSocket creation itself failed
      // cleanupConnectionInternals(); // Not strictly needed as socket is null, but harmless.
    }
  };

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

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

      if (currentStatus === "CONNECTED" && socket?.readyState === WebSocket.OPEN && currentWebSocketUrl === newWebSocketUrl) {
        console.warn(`VoiceWS (connect [${localCallId}]): Already connected to the same URL.`);
        resetConnectingFlag(); // Ensure it's false if truly connected.
        return;
      }

      // Prevent concurrent explicit `connect` calls if one is already in `CONNECTING` or `AUTHENTICATING`
      if (isConnectingInProgress && (currentStatus === "CONNECTING" || currentStatus === "AUTHENTICATING")) {
        console.warn(`VoiceWS (connect [${localCallId}]): Connect called while a connection attempt (${currentStatus}) is already in progress. Aborting new call.`);
        return;
      }

      // Reset flags for a *new* explicit connection sequence
      isIntentionalDisconnect = false;
      serverReportedError = false;
      // isConnectingInProgress will be set by _connect if it proceeds.
      // Ensure it's reset if we are cleaning up an old connection attempt first.
      resetConnectingFlag();


      if (socket) { // If there's an old socket (e.g. from a failed/interrupted attempt)
        console.warn(`VoiceWS (connect [${localCallId}]): Explicit connect called with existing socket (state: ${socket.readyState}). Cleaning up old one.`);
        cleanupConnectionInternals(); // This will close the old socket and nullify it.
      }

      currentWebSocketUrl = newWebSocketUrl;
      getTokenFunc = newGetTokenFunc;

      _connect(); // Start the connection process
    },

    disconnect: (intentional = true) => {
      const callId = currentConnectCallId;
      console.log(`VoiceWS (disconnect [${callId}]): Disconnect requested. Intentional: ${intentional}, InProgress: ${isConnectingInProgress}`);

      const previousStatus = get().status;
      const previousError = get().lastError;

      isIntentionalDisconnect = intentional;
      if (intentional) {
        // User action clears server error flag, allowing a future manual connect to proceed
        // without being blocked by a previous serverReportedError.
        serverReportedError = false;
      }
      resetConnectingFlag(); // Stop any perceived ongoing connection attempt.

      cleanupConnectionInternals(); // This will close the socket if open, triggering onclose.

      // Set final state. `handleClose` will also run and might refine `lastError`
      // based on close event details if it was an unexpected close.
      if (previousStatus === "ERROR" && !intentional) {
        // If it was already an error and this disconnect is not user-initiated (e.g. internal call)
        set({ status: "ERROR", lastError: previousError || "Disconnected due to an error." });
      } else {
        set({
          status: "DISCONNECTED",
          lastError: intentional ? "User disconnected" : (previousError || "Disconnected"),
        });
      }
    },

    sendSdpOffer: (sdp) => {
      if (get().status !== "CONNECTED") {
        console.warn("VoiceWS: Cannot send SDP_OFFER. Not connected.");
        // Optionally set lastError if this is considered an application error
        // set(state => ({ lastError: state.lastError || "Attempted to send offer while not connected." }));
        return false;
      }
      return sendWebSocketMessage({ type: "SDP_OFFER", data: { sdp } });
    },

    setCallbacks: (newCallbacks) => {
      voiceCallbacks = { ...voiceCallbacks, ...newCallbacks };
      console.debug("VoiceWS: Callbacks updated.", voiceCallbacks);
    },
  };
});

// Optional: Online/Offline listeners
if (typeof window !== 'undefined') {
  const handleBrowserOnline = () => {
    const { status } = useVoiceWebSocketStore.getState();
    console.log("VoiceWS: Browser online event detected.");

    if (['DISCONNECTED', 'ERROR'].includes(status) && !isConnectingInProgress) {
      console.log("VoiceWS: Browser is online. Application may re-attempt connection if desired by calling connect().");
      // Example: Notify the application if it wants to handle this
      // if (voiceCallbacks.onBrowserOnline) {
      //   voiceCallbacks.onBrowserOnline();
      // }
      // Or dispatch a global event:
      // window.dispatchEvent(new CustomEvent('voiceWsBrowserOnline'));
    } else {
      console.log(`VoiceWS: Browser online, no action needed. Status: ${status}, ServerError: ${serverReportedError}, ConnectingInProgress: ${isConnectingInProgress}`);
    }
  };
  window.addEventListener('online', handleBrowserOnline);

  const handleAppClose = () => { // Best-effort for page unload
    const { disconnect, status } = useVoiceWebSocketStore.getState();
    if (status !== "IDLE" && status !== "DISCONNECTED") {
      console.log("VoiceWS: Disconnecting due to page unload.");
      disconnect(true); // Intentional disconnect
    }
  };
  window.addEventListener('beforeunload', handleAppClose);
}