littlefs/src/components/WebSocketProvider.tsx

import { useEffect } from "preact/hooks";
import { $websocket } from "./WebSocketAtom";
import {
  stateSetters,
  LITTLEFS_VERSION,
  type UpdateInfo,
  type UpdateStatus,
  type UpdateTarget,
} from "./AppProviderStore";
import {
  $scanned_wifi_networks,
  $current_network_configuration,
  $network_status_display,
  $network_config_error,
  request_network_config,
  type WiFiNetwork,
  authmodeToSecurity,
  ConnectionMode,
} from "./NetworkManagerStore";
import { send_webinterface_version } from "./ota_update_components/OtaCommands";

export { $websocket } from "./WebSocketAtom";

// Heartbeat configuration
const HEARTBEAT_INTERVAL = 1500; // Send ping every 1.5 seconds
const HEARTBEAT_TIMEOUT = 1000;  // Wait 1 second for pong response
const MAX_MISSED_HEARTBEATS = 5; // Close connection after 5 missed pongs

// Heartbeat state
let heartbeatInterval: ReturnType<typeof setInterval> | null = null;
let heartbeatTimeout: ReturnType<typeof setTimeout> | null = null;
let missedHeartbeats = 0;
let _lastPingSentAt = 0;

// Reconnect state
let isReconnecting = false;
let reconnectTimeout: ReturnType<typeof setTimeout> | null = null;

function clearHeartbeat() {
  if (heartbeatInterval) {
    clearInterval(heartbeatInterval);
    heartbeatInterval = null;
  }
  if (heartbeatTimeout) {
    clearTimeout(heartbeatTimeout);
    heartbeatTimeout = null;
  }
  missedHeartbeats = 0;
}

function clearReconnect() {
  if (reconnectTimeout) {
    clearTimeout(reconnectTimeout);
    reconnectTimeout = null;
  }
  isReconnecting = false;
}

function scheduleReconnect() {
  if (isReconnecting) {
    return;
  }
  isReconnecting = true;
  reconnectTimeout = setTimeout(() => {
    connectWebSocket();
    isReconnecting = false;
  }, 1000);
}

function sendHeartbeat() {
  const ws = $websocket.get();
  if (ws && (ws as WebSocket).readyState === (ws as WebSocket).OPEN) {
    // Send ping message via WebSocket
    const pingMessage = JSON.stringify({ message_type: "ping" });
    try {
      (ws as WebSocket).send(pingMessage);
      _lastPingSentAt = performance.now();
    } catch (error) {
      missedHeartbeats++;
      return;
    }

    // Set timeout for pong response
    if (heartbeatTimeout) {
      clearTimeout(heartbeatTimeout);
      heartbeatTimeout = null;
    }
    heartbeatTimeout = setTimeout(() => {
      missedHeartbeats++;
      if (missedHeartbeats >= MAX_MISSED_HEARTBEATS) {
        // First stop the heartbeat, then close + null the socket
        clearHeartbeat();
        const cur = $websocket.get();
        if (cur) {
          (cur as WebSocket).close();
          stateSetters.setWebSocketConnectionStatus("Disconnected");
          $websocket.set(null);
        }
      }
    }, HEARTBEAT_TIMEOUT);
  }
}

function startHeartbeat() {
  clearHeartbeat(); // Clear any existing heartbeat
  heartbeatInterval = setInterval(() => {
    sendHeartbeat();
  }, HEARTBEAT_INTERVAL);
}

function onOpen() {
  stateSetters.setWebSocketConnectionStatus("Connected");
  missedHeartbeats = 0;

  // Start heartbeat mechanism
  startHeartbeat();

  // Ask the device for its current network config so the UI knows whether to show
  // the first-time wizard or the configured overview.
  setTimeout(() => {
    request_network_config();
  }, 100);

  // Inform the firmware which webinterface version this bundle was built with.
  // The firmware uses this as update_info_s3_webinterface.currentVersion for the
  // OTA check, so independent FS-only updates work correctly. See
  // set_webinterface_version handler in websocket_server.cpp.
  setTimeout(() => {
    console.log("[OTA] sending LITTLEFS_VERSION=", LITTLEFS_VERSION);
    send_webinterface_version(LITTLEFS_VERSION);
  }, 200);
}

// Map FW status strings (target_update_status_t -> getStatusString in
// update_notifier.cpp) to the web app's UpdateStatus union.
function mapStatus(s: string | undefined): UpdateStatus {
  switch (s) {
    case "no_update":
    case "started":
    case "in_progress":
    case "completed":
    case "error":
    case "restarting":
    case "waiting":
      return s;
    default:
      console.warn("[OTA] unknown status string from FW:", s);
      return "no_update";
  }
}

function isUpdateTarget(s: string | undefined): s is UpdateTarget {
  return s === "s3_firmware" || s === "s3_webinterface";
}

// Translate an `update_progress` frame into atom updates + drive the modal
// state machine. Called from onMessage.
function handleUpdateProgress(target: UpdateTarget, progress: number, status: UpdateStatus, errorMsg?: string) {
  console.log(`[OTA] progress: target=${target} progress=${progress} status=${status}`,
              errorMsg ? `error=${errorMsg}` : "");
  if (target === "s3_firmware") {
    stateSetters.setS3FirmwareProgress(progress, status);
  } else {
    stateSetters.setS3WebinterfaceProgress(progress, status);
  }
  if (status === "error") {
    if (errorMsg) stateSetters.setUpdateErrorFor(target, errorMsg);
    stateSetters.setUpdateState("error");
  } else if (status === "restarting") {
    stateSetters.setUpdateState("restarting");
  } else if (status === "started" || status === "in_progress") {
    stateSetters.setUpdateState("updating");
  } else if (status === "completed") {
    // Completion of one target doesn't necessarily mean the whole sequence is
    // done — the firmware sends per-target completed frames and a final
    // restart frame when the firmware part lands. We let device_restarting /
    // a separate aggregate trigger flip to 'success'.
    if (target === "s3_firmware") {
      // FW always restarts after its own update → next frame will be
      // device_restarting. Mark as restarting now so the success popup picks
      // up the right copy.
      stateSetters.setUpdateState("restarting");
    } else {
      // Web-only successful update: no restart, switch straight to success.
      stateSetters.setUpdateState("success");
    }
  }
}

function onMessage(message: MessageEvent) {
  // Verarbeite JSON Nachrichten
  const dataFromServer = JSON.parse(message.data.toString());

  if (dataFromServer.message_type === "pong") {
    // reset missed heartbeats and clear pending timeout
    if (heartbeatTimeout) {
      clearTimeout(heartbeatTimeout);
      heartbeatTimeout = null;
    }
    missedHeartbeats = 0;
    if ((window as any).__wsDebug && _lastPingSentAt > 0)
      console.log(`[WS] pong RTT=${(performance.now() - _lastPingSentAt).toFixed(1)}ms`);
  } else if (dataFromServer.message_type === "wifi_networks_response") {
    const networks: WiFiNetwork[] = (dataFromServer.networks ?? []).map((n: any) => ({
      SSID: n.ssid ?? "",
      RSSI: n.rssi ?? 0,
      Channel: n.channel ?? 0,
      SecurityType: authmodeToSecurity(n.authmode ?? 0),
    }));
    $scanned_wifi_networks.set(networks);
  } else if (dataFromServer.message_type === "network_status_update") {
    $network_status_display.set({
      Connection_Mode: dataFromServer.Connection_Mode ?? ConnectionMode.CONNECTION_MODE_UNKNOWN,
      AP_Channel: dataFromServer.AP_Channel ?? 0,
      IP_Address_Station: dataFromServer.IP_Address_Station ?? "",
      Gateway_IP_Address_Station: dataFromServer.Gateway_IP_Address_Station ?? "",
      Netmask_Station: dataFromServer.Netmask_Station ?? "",
      MDNS_Hostname: dataFromServer.MDNS_Hostname ?? "",
      MAC_Address_WiFi: dataFromServer.MAC_Address_WiFi ?? "",
      rssi: dataFromServer.rssi ?? null,
    });
  } else if (dataFromServer.message_type === "network_config_response") {
    console.log("[WS] network_config_response received:", dataFromServer);
    $current_network_configuration.set({
      Connection_Mode: dataFromServer.Connection_Mode ?? ConnectionMode.CONNECTION_MODE_UNKNOWN,
      DHCP_Mode_Station: dataFromServer.DHCP_Mode_Station ?? true,
      AP_Channel: dataFromServer.AP_Channel ?? 0,
      IP_Address_Station: dataFromServer.IP_Address_Station ?? "",
      Gateway_IP_Address_Station: dataFromServer.Gateway_IP_Address_Station ?? "",
      Netmask_Station: dataFromServer.Netmask_Station ?? "",
      Static_Cfg_IP_Address_Station: dataFromServer.Static_Cfg_IP_Address_Station ?? "",
      Static_Cfg_Gateway_IP_Address_Station: dataFromServer.Static_Cfg_Gateway_IP_Address_Station ?? "",
      Static_Cfg_Netmask_Station: dataFromServer.Static_Cfg_Netmask_Station ?? "",
      SSID_Station: dataFromServer.SSID_Station ?? "",
      Password_Station: dataFromServer.Password_Station ?? "",
      SSID_AP: dataFromServer.SSID_AP ?? "",
      Password_AP: dataFromServer.Password_AP ?? "",
      MDNS_Hostname: dataFromServer.MDNS_Hostname ?? "hydrobuddy",
    });
    // Update MAC in status display — config response is the first message that carries it
    if (dataFromServer.MAC_Address_WiFi) {
      $network_status_display.set({
        ...$network_status_display.get(),
        MAC_Address_WiFi: dataFromServer.MAC_Address_WiFi || $network_status_display.get().MAC_Address_WiFi,
      });
    }
    console.log("[WS] $current_network_configuration updated:", $current_network_configuration.get());
  } else if (dataFromServer.message_type === "network_config_failed") {
    const failedSsid: string = dataFromServer.failed_ssid ?? "";
    const reason: string = dataFromServer.reason ?? "unknown";
    const connectionMode: ConnectionMode = dataFromServer.connection_mode ?? ConnectionMode.CONNECTION_MODE_UNKNOWN;
    console.warn(`[WS] network_config_failed: ssid='${failedSsid}' reason='${reason}'`);
    $network_config_error.set({ reason, failed_ssid: failedSsid, connection_mode: connectionMode });
  } else if (dataFromServer.message_type === "error_popup") {
    console.warn(`[error_popup] ${dataFromServer.error ?? dataFromServer.message ?? "(no detail)"}`);
  } else if (dataFromServer.message_type === "update_information") {
    // Server-side response to FIRMWARE_UPDATE_CMD_REQUEST_AVAILABLE_UPDATES.
    // Shape: { update_available: { s3_firmware?: {...}, s3_webinterface?: {...} }, changelog }
    console.log("[OTA] update_information received", dataFromServer);
    const ua = dataFromServer.update_available ?? {};
    const updates: UpdateInfo[] = [
      {
        target: "s3_webinterface",
        label: "Webinterface",
        currentVersion: ua.s3_webinterface?.current_version ?? LITTLEFS_VERSION,
        newVersion: ua.s3_webinterface?.new_version ?? "",
        updateAvailable: !!ua.s3_webinterface?.updateAvailable,
      },
      {
        target: "s3_firmware",
        label: "S3 Firmware",
        currentVersion: ua.s3_firmware?.current_version ?? "",
        newVersion: ua.s3_firmware?.new_version ?? "",
        updateAvailable: !!ua.s3_firmware?.updateAvailable,
      },
    ];
    stateSetters.setAvailableUpdates(updates);
    stateSetters.setUpdateChangelog(dataFromServer.changelog ?? "");
    if (updates.some((u) => u.updateAvailable)) {
      console.log("[OTA] updates available — opening UpdaterModal");
      stateSetters.setUpdaterModalOpen(true);
    } else {
      console.log("[OTA] no updates available");
    }
  } else if (dataFromServer.message_type === "update_progress") {
    if (!isUpdateTarget(dataFromServer.target)) {
      console.warn("[OTA] update_progress with unknown target:", dataFromServer.target);
    } else {
      handleUpdateProgress(
        dataFromServer.target,
        Number(dataFromServer.progress ?? 0),
        mapStatus(dataFromServer.status),
        dataFromServer.error,
      );
    }
  } else if (dataFromServer.message_type === "update_in_progress") {
    console.log("[OTA] update_in_progress=", dataFromServer.in_progress, "target=", dataFromServer.target);
    stateSetters.setUpdateInProgress(!!dataFromServer.in_progress);
  } else if (dataFromServer.message_type === "device_restarting") {
    console.log("[OTA] device_restarting received — FW is rebooting");
    stateSetters.setUpdateState("restarting");
  } else {
    // Unbekannter Nachrichtentyp - korrektes Feld protokollieren
    console.warn(`[${new Date().toISOString()}] Unbekannter Nachrichtentyp: ${dataFromServer.message_type ?? "undefined"}`);
    console.debug("Payload:", dataFromServer);
  }
}

function onError(error: Event) {
  stateSetters.setWebSocketConnectionStatus("Disconnected");
  console.log("WebSocket error:", error);
  clearHeartbeat();
  scheduleReconnect();
}

function onClose(event: CloseEvent) {
  stateSetters.setWebSocketConnectionStatus("Disconnected");
  console.log("WebSocket closed with code:", event.code);
  console.log("Close reason:", event.reason);
  console.log("Clean close:", event.wasClean);
  clearHeartbeat();
  if (event.code === 1000) {
    console.log("Normal closure");
  }
  scheduleReconnect();
}

function getWsAddress(): string {
  // Derive WS address from the page's origin so it works with any hostname or IP.
  const h = window.location.hostname;
  if (h === "localhost" || h === "127.0.0.1") {
    return (import.meta.env.PUBLIC_DEV_WS_HOST as string | undefined) ?? "ws://hydrobuddy.local/ws";
  }
  return `ws://${h}/ws`;
}

const connectWebSocket = () => {
  // Clear any existing reconnect attempts
  clearReconnect();
  clearHeartbeat();

  const wsAddress = getWsAddress();
  console.log(`Verbinde mit: ${wsAddress}`);
  if ($websocket.get() !== null) {
    ($websocket.get() as WebSocket).close(1000);
    $websocket.set(null);
  }
  $websocket.set(new WebSocket(wsAddress));
  const ws = $websocket.get();
  if (ws !== null) {
    (ws as WebSocket).onopen = onOpen;
    (ws as WebSocket).onmessage = onMessage;
    (ws as WebSocket).onerror = onError;
    (ws as WebSocket).onclose = onClose;
  }
};

export const WebSocketProvider = () => {
  // Auto-connect WebSocket at first render
  useEffect(() => {
    console.log("Component mounted, starting WebSocket connection...");
    connectWebSocket();
    // Cleanup beim Unmounting
    return () => {
      clearHeartbeat();
      clearReconnect();
      const cur = $websocket.get();
      if (cur) (cur as WebSocket).close();
      $websocket.set(null);
      stateSetters.setWebSocketConnectionStatus("Disconnected");
    };
  }, []);
};