import WebSocket from "ws";
import {
  ISocketClient,
  MethodHandler,
  MethodHandlers,
  RequestContextBase,
  SocketTimeouts,
  SocketErrorException,
  GenericMiddleware,
  NonEmptyArray,
} from "../types";

interface SocketRequest<Payload = unknown> {
  method: string;
  payload: Payload;
  id: number;
  setAuthorization?: string;
}

interface SocketResponse<Payload = unknown> {
  id: number;
  payload: Payload;
  error: SocketError | null;
}

export interface SocketMessage<
  RequestPayload = unknown,
  ResponsePayload = unknown,
> {
  request?: SocketRequest<RequestPayload>;
  response?: SocketResponse<ResponsePayload>;
}

export interface SocketError {
  message: string;
  code: number;
}

/**
 * ISocket is a class that wraps a WebSocket connection and provides a simple interface for sending and receiving messages.
 * ISocket is initialized by both server and client and request/response is called symmetrically.
 *
 * ISocket handles has two timeout actions:
 * 1. Connection timeout: If the connection is not closed after a certain time, the connection is closed.
 *  - This is useful to prevent a connection from being open indefinitely. This is mainly used in the server side since
 *    the client side manages connection lifecycle manually.
 * 2. No response timeout: If the response is not received after a certain time, the request is timed out.
 *  - This is mainly used by CLI to handle the case where the server is not responding.
 *
 * It is expected that timeouts should be relatively sorted in the order of connection timeout > request timeout > no response timeout.
 */
export class ISocket<
  ImplementedMethods,
  CallableMethods,
  RequestContext extends RequestContextBase,
> {
  private readonly ws: WebSocket;
  private readonly requestHandlers: MethodHandlers<
    ImplementedMethods,
    CallableMethods,
    RequestContext
  >;
  private readonly globalMiddlewares: GenericMiddleware<
    CallableMethods,
    RequestContext
  >[];
  private readonly responseHandler = new Map<
    number,
    {
      resolve: (data: unknown) => void;
      reject: (error: SocketError) => void;
      timeout?: NodeJS.Timeout;
    }
  >();
  private peerAuthorization?: string;
  protected nxtRequestId: number;
  private connectionTimeout?: NodeJS.Timeout;
  // pino library not available in shared
  private logger: { error: (message?: string) => void };
  private timeouts?: SocketTimeouts;

  constructor(
    ws: WebSocket,
    requestHandlers: MethodHandlers<
      ImplementedMethods,
      CallableMethods,
      RequestContext
    >,
    globalMiddlewares: GenericMiddleware<CallableMethods, RequestContext>[],
    timeouts?: SocketTimeouts,
    logger?: { error: (message?: string) => void },
  ) {
    this.ws = ws;
    this.requestHandlers = requestHandlers;
    this.globalMiddlewares = globalMiddlewares;
    this.nxtRequestId = 0;
    this.logger = logger ?? { error: console.error };
    this.timeouts = timeouts;
    // reset connection timeout in each message received. It means connection still active
    this.resetConnectionTimeout();

    this.ws.addEventListener(
      "message",
      // eslint-disable-next-line @typescript-eslint/no-misused-promises
      async (event: WebSocket.MessageEvent) => {
        // eslint-disable-next-line @typescript-eslint/no-base-to-string
        const eventData: SocketMessage = JSON.parse(event.data.toString());
        return this.handleMessage(eventData);
      },
    );
  }

  protected async handleMessage(message: SocketMessage): Promise<void> {
    this.resetConnectionTimeout();
    if (message.request) {
      // Split the method string into parts
      const parts = message.request.method.split(".");
      let handlers = this.requestHandlers as Record<string, unknown>;
      for (const part of parts) {
        handlers = handlers[part] as Record<string, unknown>;
        if (!handlers) {
          return this.respondError(message.request.id, {
            code: 2,
            message: `unknown method ${message.request.method}`,
          });
        }
      }

      if (!Array.isArray(handlers)) {
        return this.respondError(message.request.id, {
          code: 2,
          message: "unknown method",
        });
      }
      handlers = [...this.globalMiddlewares, ...handlers] as MethodHandlers<
        ImplementedMethods,
        CallableMethods,
        RequestContext
      >;
      if (message.request.setAuthorization) {
        this.peerAuthorization = message.request.setAuthorization;
      }
      const reqCtx = {
        peerAuthorization: this.peerAuthorization,
        method: message.request.method,
        requestId: message.request.id,
      } as RequestContext;
      const client = createISocketClient(this);
      const payload = message.request.payload;
      let alreadyResponded = false;
      const generateNextFn = (idx: number): (() => Promise<unknown>) => {
        let wasCalled = false;
        return async () => {
          if (alreadyResponded) {
            throw new SocketErrorException(
              4,
              "next() was called after the response was sent",
            );
          }
          const handler = (
            handlers as unknown as NonEmptyArray<
              MethodHandler<unknown, unknown, CallableMethods, RequestContext>
            >
          )[idx];
          if (!handler) {
            throw new SocketErrorException(
              5,
              "cannot call past the last handler in the chain",
            );
          }
          if (wasCalled) {
            throw new SocketErrorException(
              6,
              "next() was called multiple times",
            );
          }
          wasCalled = true;
          return this.callHandler(
            handler,
            payload,
            reqCtx,
            client,
            generateNextFn(idx + 1),
          );
        };
      };
      let response: unknown;
      try {
        // call the first handler in the chain
        response = await generateNextFn(0)();
      } catch (error) {
        const socketError =
          error instanceof SocketErrorException
            ? { code: error.code, message: error.message }
            : { code: 3, message: (error as Error).toString() };
        return this.respondError(message.request.id, socketError);
      }
      this.respond(message.request.id, response);
      alreadyResponded = true;
    } else if (message.response && message.response.id) {
      const responseHandler = this.responseHandler.get(message.response.id);
      if (!responseHandler) {
        return;
      }
      if (message.response.error) {
        responseHandler.reject(message.response.error);
      }
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      responseHandler.resolve(message.response.payload as any);
      clearTimeout(responseHandler.timeout);
      this.responseHandler.delete(message.response.id);
    } else {
      return this.respondError(-1, {
        code: 3,
        message: "unknown request id",
      });
    }
  }

  protected async callHandler<
    Params,
    Result,
    CallableMethods,
    RequestContext extends RequestContextBase,
  >(
    handler: MethodHandler<Params, Result, CallableMethods, RequestContext>,
    params: Params,
    ctx: RequestContext,
    client: ISocketClient<CallableMethods>,
    next: () => Promise<Result>,
  ): Promise<Result> {
    return handler(params, ctx, client, next);
  }

  public request<Params, Result>(
    method: string,
    params: Params,
    authorization?: string,
  ): Promise<Result> {
    return new Promise<Result>((resolve, reject) => {
      const requestId = ++this.nxtRequestId;
      this.responseHandler.set(requestId, {
        resolve: (result) => resolve(result as Result),
        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
        reject: (error: SocketError) => reject(error),
      });
      const toSend: SocketMessage = {
        request: {
          method,
          payload: params,
          id: requestId,
          setAuthorization: authorization,
        },
      };
      this.ws.send(JSON.stringify(toSend));
      this.resetNoResponseTimeout(requestId);
    });
  }

  protected respond<Result>(requestId: number, result: Result): void {
    const toSend: SocketMessage = {
      response: {
        payload: result,
        id: requestId,
        error: null,
      },
    };
    return this.ws.send(JSON.stringify(toSend));
  }

  protected respondError(requestId: number, error: SocketError): void {
    const toSend: SocketMessage = {
      response: {
        payload: null,
        id: requestId,
        error: error,
      },
    };
    return this.ws.send(JSON.stringify(toSend));
  }

  private resetConnectionTimeout(): void {
    if (!this.timeouts?.connectionTimeoutInSeconds) {
      return;
    }
    if (this.connectionTimeout) {
      clearTimeout(this.connectionTimeout);
    }
    // Set a new timeout for the next message
    this.connectionTimeout = setTimeout(
      this.handleConnectionTimeout(),
      this.timeouts?.connectionTimeoutInSeconds * 1000,
    );
  }

  private handleConnectionTimeout() {
    return () => {
      this.logger.error(
        `Connection timed out after ${this.timeouts?.connectionTimeoutInSeconds} seconds`,
      );
      this.close();
    };
  }

  private resetNoResponseTimeout(requestId: number): void {
    if (!this.timeouts?.noResponseTimeoutInSeconds) {
      return;
    }
    const responseHandler = this.responseHandler.get(requestId);
    if (!responseHandler) {
      return;
    }
    const noResponseTimeout = responseHandler.timeout;
    const reject = responseHandler.reject;
    if (responseHandler.timeout) {
      clearTimeout(noResponseTimeout);
    }
    // Set a new timeout for the next message
    responseHandler.timeout = setTimeout(
      this.handleNoResponseTimeout(reject),
      this.timeouts?.noResponseTimeoutInSeconds * 1000,
    );
  }

  private handleNoResponseTimeout(reject: (error: SocketError) => void) {
    return () => {
      const message = `Request timed out after ${this.timeouts?.noResponseTimeoutInSeconds} seconds`;
      this.logger.error(message);
      reject({ code: 7, message });
    };
  }

  public close(): void {
    clearTimeout(this.connectionTimeout);
    this.responseHandler.forEach((handler, key) => {
      clearTimeout(handler.timeout);
      this.logger.error(
        `Rejecting pending requestId ${key} due to connection close`,
      );
      handler.reject({ code: 8, message: "Connection closed" });
    });
    this.ws.close();
  }
}

const proxyTarget = Object.freeze(() => {
  /* return nothing */
});

function createIsocketProxy<
  ImplementedMethods,
  CallableMethods,
  RequestContext extends RequestContextBase,
>(
  socket: ISocket<ImplementedMethods, CallableMethods, RequestContext>,
  // if path is undefined, it means the current object is the root object
  path: string | undefined,
): unknown {
  return new Proxy(proxyTarget, {
    get(_target, prop: string) {
      const childPath = path ? `${path}.${prop}` : prop;
      // sometimes, when `createISocketClient` is called from an async function, JS will implicitly call the `then` method on
      // its return value, because promises can be arbitrarily nested
      // so return undefined for the `then` method to avoid this
      if (childPath === "then") {
        return undefined;
      }
      return createIsocketProxy(socket, childPath);
    },

    apply(_target, _thisArg, args: unknown[]) {
      if (path === undefined) {
        throw new Error("The root object is not callable");
      }
      if (
        path.endsWith(".apply") &&
        args.length === 2 &&
        Array.isArray(args[1])
      ) {
        path = path.slice(0, -".apply".length);
        args = args[1];
      }
      return socket.request(path, args[0]);
    },
  });
}

export function createISocketClient<
  CallableMethods,
  ImplementedMethods,
  RequestContext extends RequestContextBase,
>(
  socket: ISocket<ImplementedMethods, CallableMethods, RequestContext>,
): ISocketClient<CallableMethods> {
  return {
    close: () => socket.close(),
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    call: createIsocketProxy(socket, undefined) as any,
  };
}