type MessageHandler = (msg: any) => void;
type StatusHandler = (connected: boolean) => void;

interface QueuedMessage {
  type: string;
  data: any;
}

export class WsClient {
  private ws: WebSocket | null = null;
  private handlers = new Map<string, Set<MessageHandler>>();
  private statusHandlers = new Set<StatusHandler>();
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
  private url: string;
  private queue: QueuedMessage[] = [];
  private intentionalClose = false;
  private reconnectDelay = 1000;
  private static MAX_RECONNECT_DELAY = 8000;
  private tokenGetter: (() => string | null) | null = null;
  // A self-hosted bot's carrier tunnel can drop the underlying half-open browser socket without
  // ever firing 'close' (the Cloudflare DO closes the agent's carrier but doesn't proactively
  // close browser-facing WS proxies it was multiplexing) — so reconnect, which is entirely
  // onclose-driven below, would otherwise never engage and this client sits dead forever (stuck
  // showing whatever state it was in, e.g. "Updating…", until a manual page refresh). Track pong
  // liveness as a backstop (mirrors supervisor/relay-tunnel.ts's PING_MS/PONG_TIMEOUT_MS on the
  // agent side) — but a self-update's real downtime is typically only 1-2s, far shorter than
  // this 30s timeout, so it alone left the chat stuck for up to a minute+ after the server was
  // already back. The liveness poll below is the fast path: a plain fetch has no "half-open"
  // ambiguity (it either reaches the new process or fails cleanly), so it notices the server is
  // back and forces a reconnect almost immediately instead of waiting on this socket to notice.
  private lastPongAt = 0;
  private static HEARTBEAT_MS = 15_000;
  private static PONG_TIMEOUT_MS = 30_000;
  private livenessTimer: ReturnType<typeof setInterval> | null = null;
  private static LIVENESS_POLL_MS = 1_500;
  // Bumped on every connect()/forceReconnect() so a superseded socket's own async onclose/onerror
  // (close() never fires them synchronously) can tell it's stale and no-op instead of firing a
  // spurious disconnect notification or scheduling a second, redundant reconnect.
  private generation = 0;

  constructor(url?: string, tokenGetter?: (() => string | null) | null) {
    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
    const host = import.meta.env.DEV ? 'localhost:7400' : location.host;
    this.url = url ?? `${proto}//${host}/ws`;
    this.tokenGetter = tokenGetter ?? null;
  }

  connect(): void {
    this.intentionalClose = false;
    this.startLivenessPoll();
    const gen = ++this.generation;
    let wsUrl = this.url;
    if (this.tokenGetter) {
      const token = this.tokenGetter();
      if (token) {
        const sep = wsUrl.includes('?') ? '&' : '?';
        wsUrl = `${wsUrl}${sep}token=${token}`;
      }
    }
    const ws = new WebSocket(wsUrl);
    this.ws = ws;

    ws.onopen = () => {
      if (gen !== this.generation) return;
      this.reconnectDelay = 1000;
      this.lastPongAt = Date.now();
      this.notifyStatus(true);
      this.flushQueue();
      this.startHeartbeat();
    };

    ws.onmessage = (e) => {
      if (gen !== this.generation) return;
      // Pong frame — liveness proof, not app data.
      if (e.data === 'pong') { this.lastPongAt = Date.now(); return; }
      let msg: any;
      try { msg = JSON.parse(e.data as string); } catch { return; }
      const handlers = this.handlers.get(msg.type);
      handlers?.forEach((h) => h(msg.data));
    };

    ws.onclose = () => {
      if (gen !== this.generation) return; // superseded by forceReconnect()/a newer connect()
      this.stopHeartbeat();
      this.notifyStatus(false);
      if (!this.intentionalClose) {
        this.reconnectTimer = setTimeout(() => {
          this.reconnectDelay = Math.min(this.reconnectDelay * 2, WsClient.MAX_RECONNECT_DELAY);
          this.connect();
        }, this.reconnectDelay);
      }
    };

    ws.onerror = () => { if (gen === this.generation) ws.close(); };
  }

  disconnect(): void {
    this.intentionalClose = true;
    this.generation++; // invalidate any in-flight socket's handlers
    if (this.reconnectTimer) { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; }
    this.stopHeartbeat();
    this.stopLivenessPoll();
    this.ws?.close();
    this.ws = null;
  }

  on(type: string, handler: MessageHandler): () => void {
    if (!this.handlers.has(type)) this.handlers.set(type, new Set());
    this.handlers.get(type)!.add(handler);
    return () => this.handlers.get(type)?.delete(handler);
  }

  onStatus(handler: StatusHandler): () => void {
    this.statusHandlers.add(handler);
    return () => this.statusHandlers.delete(handler);
  }

  send(type: string, data: any): void {
    const message = { type, data };
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
    } else {
      // Queue for delivery on reconnect
      this.queue.push(message);
    }
  }

  get connected(): boolean {
    return this.ws?.readyState === WebSocket.OPEN;
  }

  private flushQueue(): void {
    while (this.queue.length > 0 && this.ws?.readyState === WebSocket.OPEN) {
      const msg = this.queue.shift()!;
      this.ws.send(JSON.stringify(msg));
    }
  }

  private notifyStatus(connected: boolean): void {
    this.statusHandlers.forEach((h) => h(connected));
  }

  private startHeartbeat(): void {
    this.stopHeartbeat();
    this.lastPongAt = Date.now();
    this.heartbeatTimer = setInterval(() => {
      if (this.ws?.readyState !== WebSocket.OPEN) return;
      if (Date.now() - this.lastPongAt > WsClient.PONG_TIMEOUT_MS) {
        // Half-open socket — force-close so onclose drives the normal reconnect loop.
        try { this.ws.close(); } catch {}
        return;
      }
      this.ws.send('ping');
    }, WsClient.HEARTBEAT_MS);
  }

  private stopHeartbeat(): void {
    if (this.heartbeatTimer) { clearInterval(this.heartbeatTimer); this.heartbeatTimer = null; }
  }

  private startLivenessPoll(): void {
    if (this.livenessTimer) return;
    this.livenessTimer = setInterval(() => {
      const s = this.ws?.readyState;
      // Skip while a handshake is already in flight, or while OPEN with a pong recent enough to
      // trust — nothing to check or fix right now.
      if (s === WebSocket.CONNECTING) return;
      if (s === WebSocket.OPEN && Date.now() - this.lastPongAt < WsClient.HEARTBEAT_MS * 2) return;
      fetch('/__bloby/version', { cache: 'no-store' })
        .then((r) => { if (r.ok) this.forceReconnect(); })
        .catch(() => {});
    }, WsClient.LIVENESS_POLL_MS);
  }

  private stopLivenessPoll(): void {
    if (this.livenessTimer) { clearInterval(this.livenessTimer); this.livenessTimer = null; }
  }

  private forceReconnect(): void {
    // close() never fires onclose synchronously, and connect() below bumps `generation` before
    // that eventually happens — so the stale socket's own onclose sees itself superseded and
    // no-ops instead of double-scheduling a reconnect or firing a spurious disconnect.
    if (this.reconnectTimer) { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; }
    try { this.ws?.close(); } catch {}
    this.reconnectDelay = 1000;
    this.connect();
  }
}