import { Injectable } from '@nestjs/common';
import type { WebSocket } from 'ws';

@Injectable()
export class RealtimeService {
  private readonly channels = new Map<string, Set<WebSocket>>();

  subscribe(channel: string, socket: WebSocket) {
    let sockets = this.channels.get(channel);
    if (!sockets) {
      sockets = new Set();
      this.channels.set(channel, sockets);
    }
    sockets.add(socket);
  }

  unsubscribe(channel: string, socket: WebSocket) {
    const sockets = this.channels.get(channel);
    sockets?.delete(socket);
    if (sockets?.size === 0) {
      this.channels.delete(channel);
    }
  }

  removeSocket(socket: WebSocket) {
    for (const [channel, sockets] of this.channels) {
      sockets.delete(socket);
      if (sockets.size === 0) {
        this.channels.delete(channel);
      }
    }
  }

  publish(channel: string, message: object) {
    const sockets = this.channels.get(channel);
    if (!sockets) {
      return;
    }
    const data = JSON.stringify(message);
    for (const socket of sockets) {
      if (socket.readyState === socket.OPEN) {
        socket.send(data);
      }
    }
  }
}
