import WebSocket from "ws";
import { backoffMs } from "./ratelimit.js";

export interface ManagedWebSocketOptions {
  /** Called on every (re)connect; send subscriptions here. */
  onOpen?: (ws: ManagedWebSocket) => void;
  onMessage: (data: string, ws: ManagedWebSocket) => void;
  onClose?: (code: number, reason: string) => void;
  onError?: (err: Error) => void;
  /** Consider the connection stale after this many ms without any frame (default 60 s). */
  staleAfterMs?: number;
  /** Ping payload sent when idle (WebSocket ping frame is used when null). */
  heartbeat?: { intervalMs: number; message?: string } | null;
  headers?: Record<string, string>;
  maxBackoffMs?: number;
  label?: string;
}

type Listener = (event: "connected" | "reconnecting" | "closed" | "stale", detail?: Record<string, unknown>) => void;

/**
 * Reconnecting WebSocket with heartbeat and stale detection.
 * One connection may carry many instruments — never open one socket per symbol.
 */
export class ManagedWebSocket {
  private ws: WebSocket | null = null;
  private attempt = 0;
  private closedByUser = false;
  private lastFrameAt = 0;
  private staleTimer: NodeJS.Timeout | null = null;
  private heartbeatTimer: NodeJS.Timeout | null = null;
  private reconnectTimer: NodeJS.Timeout | null = null;
  private listeners: Listener[] = [];
  reconnects = 0;
  connectedAt: number | null = null;

  constructor(
    public readonly url: string,
    private opts: ManagedWebSocketOptions,
  ) {}

  on(listener: Listener) {
    this.listeners.push(listener);
  }

  private emit(...args: Parameters<Listener>) {
    for (const l of this.listeners) l(...args);
  }

  connect(): void {
    this.closedByUser = false;
    this.open();
  }

  private open() {
    if (this.closedByUser) return;
    const ws = new WebSocket(this.url, { headers: this.opts.headers, handshakeTimeout: 15_000 });
    this.ws = ws;
    ws.on("open", () => {
      this.attempt = 0;
      this.connectedAt = Date.now();
      this.lastFrameAt = Date.now();
      this.emit("connected");
      this.armStale();
      this.armHeartbeat();
      this.opts.onOpen?.(this);
    });
    ws.on("message", (data) => {
      this.lastFrameAt = Date.now();
      try {
        this.opts.onMessage(typeof data === "string" ? data : data.toString("utf8"), this);
      } catch (err) {
        this.opts.onError?.(err instanceof Error ? err : new Error(String(err)));
      }
    });
    ws.on("pong", () => (this.lastFrameAt = Date.now()));
    ws.on("error", (err) => this.opts.onError?.(err));
    ws.on("close", (code, reason) => {
      this.clearTimers();
      this.ws = null;
      this.opts.onClose?.(code, reason.toString());
      if (!this.closedByUser) this.scheduleReconnect();
      else this.emit("closed");
    });
  }

  private scheduleReconnect() {
    if (this.reconnectTimer) return;
    const wait = backoffMs(this.attempt, 1000, this.opts.maxBackoffMs ?? 120_000);
    this.attempt++;
    this.reconnects++;
    this.emit("reconnecting", { attempt: this.attempt, waitMs: wait });
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.open();
    }, wait);
  }

  private armStale() {
    const staleAfter = this.opts.staleAfterMs ?? 60_000;
    if (this.staleTimer) clearInterval(this.staleTimer);
    this.staleTimer = setInterval(() => {
      if (Date.now() - this.lastFrameAt > staleAfter) {
        this.emit("stale", { silentMs: Date.now() - this.lastFrameAt });
        this.ws?.terminate();
      }
    }, Math.max(5000, Math.floor(staleAfter / 3)));
  }

  private armHeartbeat() {
    const hb = this.opts.heartbeat;
    if (hb === null) return;
    const intervalMs = hb?.intervalMs ?? 25_000;
    this.heartbeatTimer = setInterval(() => {
      if (this.ws?.readyState !== WebSocket.OPEN) return;
      if (hb?.message) this.ws.send(hb.message);
      else this.ws.ping();
    }, intervalMs);
  }

  private clearTimers() {
    if (this.staleTimer) clearInterval(this.staleTimer);
    if (this.heartbeatTimer) clearInterval(this.heartbeatTimer);
    this.staleTimer = null;
    this.heartbeatTimer = null;
  }

  send(data: string | object): boolean {
    if (this.ws?.readyState !== WebSocket.OPEN) return false;
    this.ws.send(typeof data === "string" ? data : JSON.stringify(data));
    return true;
  }

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

  close(): void {
    this.closedByUser = true;
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
    this.reconnectTimer = null;
    this.clearTimers();
    this.ws?.close();
    this.ws = null;
  }
}