import { wsConfig, type SocketKey } from "../app/config/wsconfig";

type Listener = (msg: any) => void;

class WebSocketManager {
  private sockets = new Map<SocketKey, WebSocket>();
  private listeners = new Map<SocketKey, Set<Listener>>();

  private ensureSocket(key: SocketKey): WebSocket {
    if (!this.sockets.has(key)) {
      const url = wsConfig[key];
      const ws = new WebSocket(url);
      this.sockets.set(key, ws);

      ws.onmessage = async (event) => {
        const text = await event.data.text();
        const data = JSON.parse(text);

        const ls = this.listeners.get(key);

        ls?.forEach((cb) => cb(data));
      };

      ws.onclose = () => {
        // optional: handle reconnect logic here later
      };
    }
    return this.sockets.get(key)!;
  }

  subscribe(key: SocketKey, listener: Listener) {
    if (!this.listeners.has(key)) {
      this.listeners.set(key, new Set());
    }
    this.listeners.get(key)!.add(listener);

    // Ensure socket is created and listening
    this.ensureSocket(key);

    // Return unsubscribe function
    return () => {
      const set = this.listeners.get(key);
      if (!set) return;
      set.delete(listener);
    };
  }

  send(key: SocketKey, payload: any) {
    const socket = this.ensureSocket(key);
    if (socket.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify(payload));
    }
    // if not open, you might buffer / retry etc.
  }
}

export const wsManager = new WebSocketManager();