import {
  safeParseServerMessage,
  type ClientMessage,
  type ServerMessage,
} from "@agent-client/shared";

export type WsStatus = "connecting" | "open" | "closed" | "error";

export interface WsClientHandlers {
  onMessage: (msg: ServerMessage) => void;
  onStatus: (status: WsStatus, info?: string) => void;
}

export interface WsClient {
  send: (msg: ClientMessage) => void;
  close: () => void;
}

/**
 * Minimal WebSocket client with auto-reconnect and exponential backoff.
 * One connection per browser tab; lifetime tied to caller (call close on unmount).
 */
export function createWsClient(
  url: string,
  token: string,
  handlers: WsClientHandlers,
): WsClient {
  let ws: WebSocket | null = null;
  let closedByUser = false;
  let reconnectDelay = 500;

  const fullUrl = `${url}?token=${encodeURIComponent(token)}`;

  function connect() {
    handlers.onStatus("connecting");
    try {
      ws = new WebSocket(fullUrl);
    } catch (err) {
      handlers.onStatus("error", String(err));
      scheduleReconnect();
      return;
    }

    ws.onopen = () => {
      reconnectDelay = 500;
      handlers.onStatus("open");
    };

    ws.onmessage = (ev) => {
      let raw: unknown;
      try {
        raw = JSON.parse(typeof ev.data === "string" ? ev.data : "");
      } catch {
        return;
      }
      const result = safeParseServerMessage(raw);
      if (result.success) {
        handlers.onMessage(result.data);
      }
    };

    ws.onerror = () => {
      handlers.onStatus("error");
    };

    ws.onclose = (ev) => {
      handlers.onStatus("closed", `code=${ev.code}`);
      if (!closedByUser) scheduleReconnect();
    };
  }

  function scheduleReconnect() {
    if (closedByUser) return;
    const delay = reconnectDelay;
    reconnectDelay = Math.min(reconnectDelay * 2, 10_000);
    setTimeout(() => {
      if (!closedByUser) connect();
    }, delay);
  }

  connect();

  return {
    send(msg) {
      if (ws && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify(msg));
      }
    },
    close() {
      closedByUser = true;
      ws?.close();
    },
  };
}
