import { Duplex } from "stream";
import { WebSocketServer } from "ws";

export interface NextWebsocketConfig {
  app: NextWebsocketApp;
}

export interface WsRouteHandler {
  upgrade?: (req: Request, socket: Duplex, head: NonSharedBuffer) => void;
  post?: (req: Request) => void;
  get?: (req: Request) => void;
  put?: (req: Request) => void;
  delete?: (req: Request) => void;
  patch?: (req: Request) => void;
}

export class NextWebsocketApp {
  handlers: Record<string, WsRouteHandler> = {};
  websocketHandlers: ((ws: WebSocketServer) => void)[] = [];

  on<K extends keyof WsRouteHandler>(
    method: K,
    route: string,
    handler: WsRouteHandler[K],
  ) {
    if (!this.handlers[route]) this.handlers[route] = {};
    this.handlers[route][method] = handler;
  }

  wss(handler: (ws: WebSocketServer) => void) {
    this.websocketHandlers.push(handler);
  }

  upgrade(route: string, handler: WsRouteHandler["upgrade"]) {
    this.on("upgrade", route, handler);
  }

  post(route: string, handler: WsRouteHandler["post"]) {
    this.on("post", route, handler);
  }

  get(route: string, handler: WsRouteHandler["get"]) {
    this.on("get", route, handler);
  }

  emit<K extends keyof WsRouteHandler>(
    method: K,
    route: string,
    ...args: Parameters<NonNullable<WsRouteHandler[K]>>
  ) {
    const handler = this.handlers[route]?.[method] as
      | ((...args: Parameters<NonNullable<WsRouteHandler[K]>>) => unknown)
      | undefined;

    handler?.(...args);
  }
}