import { applyWSSHandler } from '@trpc/server/adapters/ws';
import { WebSocketServer } from 'ws';
import type { WebSocket } from 'ws';
import yargs from 'yargs';
import type { AppRouter, ParsedCommandOptions } from './trpc';
import { createContext } from './trpc';
import { logger } from './logger';
import { clearAllMappedFunctions } from './fnMap';
import { clearAllMappedPromises } from './promiseMap';
import { appRouter } from './router';
import { clearClientFunctionCallEventSubjects, clientFunctionReturnEventSubject } from './subjects';
import { subscribeClientFunctionCalls } from './utils';

declare module 'ws' {
  interface WebSocket {
    socketId: number;
  }
}

const params = yargs(process.argv.slice(2)).argv;

const port = (params as ParsedCommandOptions).port as number;

const logLevel = ((params as ParsedCommandOptions).logLevel as string) || 'INFO';

logLevel && (logger.level = logLevel);

logger.debug(`command line arguments: ${JSON.stringify(params)}`);

let uniqueSocketID = 1;

function startServer(port: number) {
  subscribeClientFunctionCalls();
  const wss = new WebSocketServer({ port });
  applyWSSHandler<AppRouter>({
    wss,
    router: appRouter,
    createContext
  });

  wss.on('connection', (ws: WebSocket) => {
    ws.socketId = uniqueSocketID++;
    logger.info(`➕➕ Connection (${wss.clients.size})`);
    ws.once('close', () => {
      logger.info(`➖➖ Connection (${wss.clients.size})`);
    });
  });

  logger.info(`✅ WebSocket Server listening on ws://localhost:${port}`);

  const handleExit = (code: number, error?: Error) => {
    logger.info(`exit code: ${code}`);
    wss.close();
    clearAllMappedFunctions();
    clearAllMappedPromises();
    clearClientFunctionCallEventSubjects();
    clientFunctionReturnEventSubject.unsubscribe();

    if (error) {
      logger.error(`error for process ${process.pid}: ${error}`);
    }
    process.exit(code);
  };

  process.on('exit', (code) => handleExit(code));
  // According to the specifications of POSIX, using 128 + SignalCode to be exit code.
  // Some exit codes.
  process.on('SIGHUP', () => handleExit(128 + 1));
  process.on('SIGINT', () => handleExit(128 + 2));
  process.on('SIGTERM', () => handleExit(128 + 15));
  // In windows, Pressed ctrl+break to exit.
  process.on('SIGBREAK', () => handleExit(128 + 21));
  // Exit Code 1 indicates uncaught or unhandled exception.
  process.on('uncaughtException', (error: Error) => handleExit(1, error));
  process.on('unhandledRejection', (error: Error) => handleExit(1, error));
}

startServer(port);