import { nanoid } from 'nanoid';
import { WebSocket } from 'ws';
import { logger } from '../logger.ts';
class WsMessage {
  ws: WebSocket;
  user?: string;
  constructor({ ws, user }: WssMessageOptions) {
    this.ws = ws;
    this.user = user;
  }
  async sendData(data: any, opts?: { timeout?: number }) {
    if (this.ws.readyState !== WebSocket.OPEN) {
      return { code: 500, message: 'WebSocket is not open' };
    }
    const timeout = opts?.timeout || 10 * 6 * 1000; // 10 minutes
    const id = nanoid();
    const message = JSON.stringify({
      id,
      type: 'proxy',
      data,
    });
    logger.info('ws-proxy sendData', message);
    this.ws.send(message);
    return new Promise((resolve) => {
      const timer = setTimeout(() => {
        resolve({
          code: 500,
          message: 'timeout',
        });
      }, timeout);
      this.ws.once('message', (event: Buffer) => {
        const eventData = event.toString();
        if (!eventData) {
          return;
        }
        const data = JSON.parse(eventData);
        if (data.id === id) {
          resolve(data.data);
          clearTimeout(timer);
        }
      });
    });
  }
}
type WssMessageOptions = {
  ws: WebSocket;
  user?: string;
};
export class WsProxyManager {
  wssMap: Map<string, WsMessage> = new Map();
  constructor() {}
  getId(id: string, user?: string) {
    return id + '/' + user;
  }
  register(id: string, opts?: { ws: WebSocket; user: string }) {
    const _id = this.getId(id, opts?.user || '');
    if (this.wssMap.has(_id)) {
      const value = this.wssMap.get(_id);
      if (value) {
        value.ws.close();
      }
    }
    const value = new WsMessage({ ws: opts?.ws, user: opts?.user });
    this.wssMap.set(_id, value);
  }
  unregister(id: string, user?: string) {
    const _id = this.getId(id, user || '');
    const value = this.wssMap.get(_id);
    if (value) {
      value.ws.close();
    }
    this.wssMap.delete(_id);
  }
  getIds() {
    return Array.from(this.wssMap.keys());
  }
  get(id: string, user?: string) {
    if (user) {
      const _id = this.getId(id, user);
      return this.wssMap.get(_id);
    }
    return this.wssMap.get(id);
  }
}