import type { ServerFrame, ClientFrame, EventFrame } from "./protocol";

export type WsState = "connecting" | "open" | "closed";

export interface WsManager {
  onState(cb: (state: WsState) => void): () => void;
  onFrame(cb: (frame: ServerFrame) => void): () => void;
  onAuthError(cb: () => void): () => void;
  subscribe(sessionId: string): void;
  unsubscribe(sessionId: string): void;
  /** Resubscribe all live subscriptions; call on (re)open. */
  resubscribeAll(): void;
  /** SessionIds subscribed right now (used by onOpen hooks to refetch). */
  subscriptions(): ReadonlySet<string>;
  close(): void;
}

const BACKOFF_INITIAL_MS: number = 500;
const BACKOFF_MAX_MS: number = 30_000;
const CLOSED_BY_USER: number = 4900;

export function createWsManager(url: string | (() => string)): WsManager {
  const urlFn: () => string = typeof url === "string" ? () => url : url;

  let ws: WebSocket | null = null;
  let closed = false;
  let attempt: number = 0;
  let reconnectTimer: number | undefined;
  let state: WsState = "connecting";

  const stateCbs = new Set<(s: WsState) => void>();
  const frameCbs = new Set<(f: ServerFrame) => void>();
  const authErrorCbs = new Set<() => void>();
  const subs = new Set<string>();

  const setState = (s: WsState): void => {
    state = s;
    stateCbs.forEach((cb) => cb(s));
  };

  const connect = (): void => {
    if (closed) return;
    setState("connecting");
    const sock = new WebSocket(urlFn());
    ws = sock;
    sock.onopen = () => {
      attempt = 0;
      setState("open");
      // resubscribe current chat + sessions fan-out resumes automatically
      subs.forEach((id) => {
        const frame: ClientFrame = { type: "subscribe", sessionId: id };
        sock.send(JSON.stringify(frame));
      });
    };
    sock.onmessage = (ev: MessageEvent) => {
      let frame: ServerFrame;
      try {
        frame = JSON.parse(typeof ev.data === "string" ? ev.data : "") as ServerFrame;
      } catch {
        return;
      }
      if (frame === null || typeof frame !== "object" || typeof frame.type !== "string") return;
      frameCbs.forEach((cb) => cb(frame));
    };
    sock.onclose = (ev: CloseEvent) => {
      ws = null;
      setState("closed");
      if (closed) return;
      if (ev.code === 1008) {
        // policy violation: token rejected
        authErrorCbs.forEach((cb) => cb());
        return;
      }
      const jitter: number = Math.random();
      const exp: number = Math.min(BACKOFF_INITIAL_MS * 2 ** attempt, BACKOFF_MAX_MS);
      attempt += 1;
      reconnectTimer = window.setTimeout(connect, exp * (0.5 + 0.5 * jitter));
    };
    sock.onerror = () => {
      /* onclose follows */
    };
  };

  connect();

  return {
    onState(cb) {
      stateCbs.add(cb);
      cb(state);
      return () => stateCbs.delete(cb);
    },
    onFrame(cb) {
      frameCbs.add(cb);
      return () => frameCbs.delete(cb);
    },
    onAuthError(cb) {
      authErrorCbs.add(cb);
      return () => authErrorCbs.delete(cb);
    },
    subscribe(sessionId) {
      subs.add(sessionId);
      if (ws !== null && ws.readyState === WebSocket.OPEN) {
        const frame: ClientFrame = { type: "subscribe", sessionId };
        ws.send(JSON.stringify(frame));
      }
    },
    unsubscribe(sessionId) {
      subs.delete(sessionId);
      if (ws !== null && ws.readyState === WebSocket.OPEN) {
        const frame: ClientFrame = { type: "unsubscribe", sessionId };
        ws.send(JSON.stringify(frame));
      }
    },
    resubscribeAll() {
      const sock: WebSocket | null = ws;
      if (sock !== null && sock.readyState === WebSocket.OPEN) {
        subs.forEach((id) => {
          const frame: ClientFrame = { type: "subscribe", sessionId: id };
          sock.send(JSON.stringify(frame));
        });
      }
    },
    subscriptions() {
      return subs;
    },
    close() {
      closed = true;
      if (reconnectTimer !== undefined) window.clearTimeout(reconnectTimer);
      if (ws !== null) {
        ws.onclose = null;
        ws.onmessage = null;
        ws.onerror = null;
        ws.onopen = null;
        ws.close(CLOSED_BY_USER);
        ws = null;
      }
      setState("closed");
    },
  };
}

export type { EventFrame };