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;

  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;
    let wsUrl = this.url;
    if (this.tokenGetter) {
      const token = this.tokenGetter();
      if (token) {
        const sep = wsUrl.includes('?') ? '&' : '?';
        wsUrl = `${wsUrl}${sep}token=${token}`;
      }
    }
    this.ws = new WebSocket(wsUrl);

    this.ws.onopen = () => {
      this.reconnectDelay = 1000;
      this.notifyStatus(true);
      this.flushQueue();
      this.startHeartbeat();
    };

    this.ws.onmessage = (e) => {
      // Ignore pong frames
      if (e.data === 'pong') 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));
    };

    this.ws.onclose = () => {
      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);
      }
    };

    this.ws.onerror = () => this.ws?.close();
  }

  disconnect(): void {
    this.intentionalClose = true;
    if (this.reconnectTimer) { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; }
    this.stopHeartbeat();
    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.heartbeatTimer = setInterval(() => {
      if (this.ws?.readyState === WebSocket.OPEN) {
        this.ws.send('ping');
      }
    }, 25_000);
  }

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