function reconnectWebSocket() {
  if (reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
    reconnectAttempts++;
    updateConnectionStatus(
      "connecting",
      `Reconnecting... (${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS})`,
    );

    reconnectTimeout = setTimeout(() => {
      console.log(`Attempting to reconnect... (${reconnectAttempts})`);
      connectWebSocket();
    }, RECONNECT_DELAY);
  } else {
    updateConnectionStatus(
      "disconnected",
      "Failed to connect. Please refresh the page.",
    );
    showAlert("Connection lost. Please refresh the page.", "error");
  }
}

function clearReconnection() {
  if (reconnectTimeout) {
    clearTimeout(reconnectTimeout);
    reconnectTimeout = null;
  }
  reconnectAttempts = 0;
}

function updateConnectionStatus(status, message = "") {
  const statusElement = document.getElementById("connectionStatus");
  statusElement.textContent = message || status;
  statusElement.className = `connection-status ${status}`;
  statusElement.style.display = "block";

  if (status === "connected") {
    setTimeout(() => {
      statusElement.style.display = "none";
    }, 3000);
  }
}

function connectWebSocket() {
  // Clear any existing reconnection attempts
  clearReconnection();

  const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
  const wsUrl = `${protocol}//${window.location.host}/ws`;

  updateConnectionStatus("connecting", "Connecting...");

  ws = new WebSocket(wsUrl);

  ws.onopen = () => {
    console.log("WebSocket connected");
    updateConnectionStatus("connected", "Connected");
    clearReconnection();

    ws.send(
      JSON.stringify({
        type: "auth",
        token: currentToken,
      }),
    );
    addSystemMessage("Connected to chat server");
  };

  ws.onmessage = (event) => {
    console.log("Received message:", event.data);
    try {
      const data = JSON.parse(event.data);

      if (data.type === "error") {
        showAlert(data.content, "error");
        if (
          data.content.includes("Invalid") ||
          data.content.includes("expired") ||
          data.content.includes("auth")
        ) {
          logout();
        }
      } else if (data.type === "system") {
        addSystemMessage(data.content);
      } else if (data.type === "message_deleted") {
        // Handle message deletion from server
        const messageElement = document.querySelector(
          `[data-message-id="${data.message_id}"]`,
        );
        if (messageElement) {
          messageElement.remove();
        }

        // Update chat history
        if (messageHistory[currentRecipient]) {
          messageHistory[currentRecipient] = messageHistory[
            currentRecipient
          ].filter((msg) => msg.id !== data.message_id);
        }

        // Reload chats to update last message
        loadChats();
      } else {
        handleChatMessage(data);
      }
    } catch (e) {
      console.error("Error parsing message:", e, "Raw data:", event.data);
    }
  };

  ws.onclose = (event) => {
    console.log("WebSocket disconnected:", event.code, event.reason);
    updateConnectionStatus("disconnected", "Disconnected");
    addSystemMessage("Disconnected from server");
    document.getElementById("sendButton").disabled = true;
    document.getElementById("messageInput").disabled = true;

    // Attempt to reconnect if this wasn't a normal closure
    if (event.code !== 1000 && event.code !== 1001) {
      reconnectWebSocket();
    }
  };

  ws.onerror = (error) => {
    console.error("WebSocket error:", error);
    updateConnectionStatus("disconnected", "Connection error");
    addSystemMessage("Connection error");
  };
}