import { ref, type Ref } from "vue";
import { apiClient } from "@/api/client";
import type { WsServerMessage } from "@/api/types";

/**
 * Resolve which conversation an incoming WS message belongs to.
 *
 * The backend protocol currently does NOT tag server→client messages with
 * `conversation_id`, so we route by recency: pick the most recently used
 * pending conversation. This is a heuristic — it works as long as users
 * don't fire two requests in quick succession across conversations. We
 * bias toward the *current* view if it's still pending, which is what the
 * user is watching right now.
 *
 * Exported as a pure function so it can be unit-tested without Pinia.
 */
export function resolveIncomingConvId(
  currentConversationId: string | null,
  pendingConversations: Set<string>,
  pendingLastUsedAt: Map<string, number>,
): string {
  if (currentConversationId && pendingConversations.has(currentConversationId)) {
    return currentConversationId;
  }
  // Fall back to the most recently used pending conversation.
  let best: string | null = null;
  let bestTs = 0;
  pendingLastUsedAt.forEach((ts, id) => {
    if (pendingConversations.has(id) && ts > bestTs) {
      best = id;
      bestTs = ts;
    }
  });
  return best ?? currentConversationId ?? "";
}

export interface ChatSocketOptions {
  currentConversationId: Ref<string | null>;
  pendingConversations: Ref<Set<string>>;
  pendingLastUsedAt: Ref<Map<string, number>>;
  /** Invoked for each parsed server message (→ dispatchWsEvent). */
  onMessage: (data: WsServerMessage) => void;
  /** Invoked after the socket reopens (→ _recoverTaskAfterReconnect). */
  onReconnect: () => void | Promise<void>;
  /** Invoked when the socket closes (→ clear stream-side state). */
  onDisconnect: () => void;
}

/**
 * WebSocket lifecycle composable: connection, 30s heartbeat, and 3s
 * auto-reconnect guarded by `_intentionalDisconnect` to prevent cascading
 * reconnects after an explicit `disconnectWebSocket()`.
 */
export function useChatSocket(options: ChatSocketOptions) {
  const isWsConnected = ref(false);
  const ws = ref<WebSocket | null>(null);
  let _heartbeatTimer: ReturnType<typeof setInterval> | null = null;
  let _reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  let _intentionalDisconnect = false;

  function connectWebSocket(): void {
    // Problem 6: also skip if already CONNECTING to avoid orphan sockets
    if (
      ws.value &&
      (ws.value.readyState === WebSocket.OPEN ||
        ws.value.readyState === WebSocket.CONNECTING)
    ) {
      return;
    }

    _intentionalDisconnect = false;
    const socket = apiClient.createWebSocket();

    socket.onopen = () => {
      isWsConnected.value = true;
      console.log("WebSocket connected");
      // Start heartbeat: send ping every 30s to keep connection alive
      if (_heartbeatTimer) clearInterval(_heartbeatTimer);
      _heartbeatTimer = setInterval(() => {
        if (ws.value && ws.value.readyState === WebSocket.OPEN) {
          ws.value.send(JSON.stringify({ type: "ping" }));
        }
      }, 30000);
      // Check for running tasks to resume after reconnection
      void options.onReconnect();
    };

    socket.onmessage = (event: MessageEvent) => {
      try {
        const data = JSON.parse(event.data as string) as WsServerMessage;
        console.log("[Chat WS] Received:", data.type, data);
        options.onMessage(data);
      } catch (error) {
        console.error("Failed to parse WebSocket message:", error);
      }
    };

    socket.onclose = () => {
      isWsConnected.value = false;
      // P2 #21 fix: clear per-conversation pending state to prevent stuck
      // loading state during disconnect. onReconnect will re-mark
      // conversations pending if an active task is found.
      options.pendingConversations.value = new Set();
      options.pendingLastUsedAt.value = new Map();
      // Notify stream side to clear stale streaming steps.
      options.onDisconnect();
      console.log("WebSocket disconnected");
      if (_heartbeatTimer) {
        clearInterval(_heartbeatTimer);
        _heartbeatTimer = null;
      }
      // Problem 1: do not auto-reconnect after an intentional disconnect
      if (_intentionalDisconnect) {
        return;
      }
      // Auto reconnect after 3 seconds
      if (_reconnectTimer) clearTimeout(_reconnectTimer);
      _reconnectTimer = setTimeout(() => {
        if (!ws.value || ws.value.readyState === WebSocket.CLOSED) {
          connectWebSocket();
        }
      }, 3000);
    };

    socket.onerror = (error) => {
      console.error("WebSocket error:", error);
      isWsConnected.value = false;
    };

    ws.value = socket;
  }

  /** Disconnect WebSocket and suppress auto-reconnect. */
  function disconnectWebSocket(): void {
    _intentionalDisconnect = true;
    if (_reconnectTimer) {
      clearTimeout(_reconnectTimer);
      _reconnectTimer = null;
    }
    if (_heartbeatTimer) {
      clearInterval(_heartbeatTimer);
      _heartbeatTimer = null;
    }
    if (ws.value) {
      ws.value.close();
      ws.value = null;
      isWsConnected.value = false;
    }
  }

  return {
    isWsConnected,
    ws,
    connectWebSocket,
    disconnectWebSocket,
    // Bound resolver using the latest option refs (for dispatchWsEvent ctx).
    // Arrow form avoids shadowing the exported pure function above.
    resolveIncomingConvId: () =>
      resolveIncomingConvId(
        options.currentConversationId.value,
        options.pendingConversations.value,
        options.pendingLastUsedAt.value,
      ),
  };
}
