import WebSocket from 'isomorphic-ws';

import {
  SocketErrorException,
  type ISocketClient,
  type MethodHandler,
  type MethodHandlers,
  type RequestContextBase,
  type SocketTimeouts,
  type GenericMiddleware
} from './types.js';

export const REQUEST_ABORTED_SOCKET_ERROR = {
  code: 9,
  message: '[internal] Request aborted',
  /** Distinguishes a locally synthesized abort from a server SocketError with code 9. */
  aborted: true as const
} as const;

interface SocketRequest<Payload = unknown> {
  method: string;
  payload: Payload;
  id: number;
  setAuthorization?: string;
  traceparent?: string;
  tracestate?: 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;
  stack?: string;
}

// Codes in the 4000s are left available by IETF
// https://datatracker.ietf.org/doc/html/rfc6455#section-7.4.2
export const INTENTIONAL_CLOSE_CODE = 4001;
const CONNECTION_CLOSED_MESSAGE = '[internal] Connection closed';

/**
 * 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;
      abortCleanup?: () => void;
    }
  >();
  private peerAuthorization?: string;
  protected nxtRequestId: number;
  private connectionTimeout: NodeJS.Timeout;
  // pino library not available in shared
  protected logger: { error: (message?: string) => void; info: (message?: string) => void; warn: (message?: string) => void };
  private timeouts?: SocketTimeouts;
  private isClosed = false;
  private readonly getAbortSignal?: () => AbortSignal | undefined;

  constructor(
    ws: WebSocket,
    requestHandlers: MethodHandlers<ImplementedMethods, CallableMethods, RequestContext>,
    globalMiddlewares: GenericMiddleware<CallableMethods, RequestContext>[],
    options: {
      /**
       * This callback is called when the WebSocket connection is closed,
       * either with a close handshake or immediately if the connection is dropped
       *
       * If code is INTENTIONAL_CLOSE_CODE, it was closed due to an internal timeout.
       * Otherwise, the other side closed the connection.
       *
       * @param event - The WebSocket.CloseEvent object.
       */
      onClose: (event: WebSocket.CloseEvent) => void;
      timeouts?: SocketTimeouts;
      logger?: { error: (message?: string) => void; info: (message?: string) => void; warn: (message?: string) => void };
      /** Read at request time so server-side callers can supply a per-call signal without Node ALS in this package. */
      getAbortSignal?: () => AbortSignal | undefined;
    }
  ) {
    this.ws = ws;
    this.getAbortSignal = options.getAbortSignal;
    this.requestHandlers = requestHandlers;
    this.globalMiddlewares = globalMiddlewares;
    this.nxtRequestId = 0;
    this.logger = options.logger ?? {
      error: console.error,
      warn: console.warn,
      info: console.info
    };
    this.timeouts = options.timeouts;
    // reset connection timeout in each message received. It means connection still active
    this.resetConnectionTimeout();
    this.ws.addEventListener('message', (event: WebSocket.MessageEvent) => {
      try {
        const eventData: SocketMessage = JSON.parse(event.data.toString());
        void this.handleMessage(eventData);
      } catch (error) {
        this.logger.error(`Error parsing message: ${error}`);
      }
    });

    this.ws.addEventListener('close', (event: WebSocket.CloseEvent) => {
      this.isClosed = true;

      if (this.connectionTimeout) {
        clearTimeout(this.connectionTimeout);
      }

      if (event.code !== INTENTIONAL_CLOSE_CODE) {
        // Reject any pending responses that happen after the close handshake,
        // which shouldn't be possible
        this.responseHandler.forEach((handler, key) => {
          clearTimeout(handler.timeout);
          handler.abortCleanup?.();
          this.logger.error(`Rejecting pending requestId ${key} due to connection close`);
          handler.reject({ code: 8, message: '[internal] Connection closed' });
        });
        this.responseHandler.clear();
      }
      options.onClose(event);
    });
  }

  protected async handleMessage(message: SocketMessage): Promise<void> {
    if (this.isClosed) {
      this.logger.error('Connection closed, ignoring message');
      return;
    }
    this.resetConnectionTimeout();
    if (message.request) {
      // Split the method string into parts
      const parts = message.request.method.split('.');
      let handlers = this.requestHandlers;
      for (const part of parts) {
        handlers = handlers[part];
        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[idx] as MethodHandler<unknown, unknown, CallableMethods, RequestContext> | undefined;
          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, stack: error.stack }
            : { code: 3, message: error.toString(), stack: error instanceof Error ? error.stack : undefined };
        this.logger.error(JSON.stringify(socketError));
        return this.respondError(message.request.id, socketError, error);
      }
      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);
        this.clearPendingRequestState(message.response.id);
        return;
      }
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      responseHandler.resolve(message.response.payload as any);
      this.clearPendingRequestState(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> {
    if (this.isClosed) {
      // Sockets normally will silently accept requests after the close handshake,
      // but we want to reject them to avoid confusion
      return Promise.reject(new SocketErrorException(10, '[internal] Connection closed'));
    }
    return new Promise<Result>((resolve, reject) => {
      const requestId = ++this.nxtRequestId;
      this.responseHandler.set(requestId, {
        resolve: (result) => resolve(result as Result),
        reject: (error: SocketError) => reject(error)
      });
      this.attachAbortListener(requestId, this.getAbortSignal?.());
      if (!this.responseHandler.has(requestId)) {
        return;
      }
      let toSend: SocketMessage = { request: { method, payload: params, id: requestId, setAuthorization: authorization } };
      toSend = this.decorateToSend(toSend);
      this.ws.send(JSON.stringify(toSend));
      this.resetConnectionTimeout();
      this.resetNoResponseTimeout(requestId);
    });
  }

  public requestWithTracing<Params, Result>(
    method: string,
    params: Params,
    tracingHeaders: { traceparent?: string; tracestate?: string },
    authorization?: string
  ): Promise<Result> {
    if (this.isClosed) {
      return Promise.reject(new SocketErrorException(10, '[internal] Connection closed'));
    }
    return new Promise<Result>((resolve, reject) => {
      const requestId = ++this.nxtRequestId;
      this.responseHandler.set(requestId, {
        resolve: (result) => resolve(result as Result),
        reject: (error: SocketError) => reject(error)
      });
      this.attachAbortListener(requestId, this.getAbortSignal?.());
      if (!this.responseHandler.has(requestId)) {
        return;
      }
      let toSend: SocketMessage = {
        request: {
          method,
          payload: params,
          id: requestId,
          setAuthorization: authorization,
          ...tracingHeaders // Explicit tracing headers
        }
      };
      toSend = this.decorateToSend(toSend);
      this.ws.send(JSON.stringify(toSend));
      this.resetConnectionTimeout();
      this.resetNoResponseTimeout(requestId);
    });
  }

  private attachAbortListener(requestId: number, abortSignal?: AbortSignal): void {
    if (!abortSignal) {
      return;
    }

    const onAbort = () => {
      this.logger.warn(`Rejecting pending requestId ${requestId} due to abort`);
      this.rejectPendingRequest(requestId, REQUEST_ABORTED_SOCKET_ERROR);
    };

    if (abortSignal.aborted) {
      onAbort();
      return;
    }

    const handler = this.responseHandler.get(requestId);
    if (!handler) {
      return;
    }

    abortSignal.addEventListener('abort', onAbort, { once: true });
    handler.abortCleanup = () => {
      abortSignal.removeEventListener('abort', onAbort);
    };
  }

  private clearPendingRequestState(requestId: number): void {
    const responseHandler = this.responseHandler.get(requestId);
    if (!responseHandler) {
      return;
    }
    clearTimeout(responseHandler.timeout);
    responseHandler.abortCleanup?.();
    this.responseHandler.delete(requestId);
  }

  private rejectPendingRequest(requestId: number, error: SocketError): void {
    const responseHandler = this.responseHandler.get(requestId);
    if (!responseHandler) {
      return;
    }
    clearTimeout(responseHandler.timeout);
    responseHandler.abortCleanup?.();
    this.responseHandler.delete(requestId);
    responseHandler.reject(error);
  }

  protected decorateToSend(message: SocketMessage): SocketMessage {
    return message;
  }

  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, exception?: Error): 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 () => {
      console.info(`[ai-service] Socket closing due to connection timeout`);
      this.close('Connection timeout');
    };
  }

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

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

  public close(reason?: string): void {
    // Starts close handshake without clearing in-progress requests
    // see https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close
    this.ws.close(INTENTIONAL_CLOSE_CODE, reason);
  }

  public closeAndRejectPendingRequests(reason?: string): void {
    this.responseHandler.forEach((handler, key) => {
      clearTimeout(handler.timeout);
      handler.abortCleanup?.();
      this.logger.error(`Rejecting pending requestId ${key} due to intentional connection close`);
      handler.reject({ code: 8, message: reason ? `${CONNECTION_CLOSED_MESSAGE}: ${reason}` : CONNECTION_CLOSED_MESSAGE });
    });
    this.responseHandler.clear();
    this.close(reason);
  }
}

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: (reason?: string) => socket.close(reason),
    closeAndRejectPendingRequests: (reason?: string) => socket.closeAndRejectPendingRequests(reason),
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    call: createIsocketProxy(socket, undefined) as any,
    requestWithTracing: <Params, Result>(
      method: string,
      params: Params,
      tracingHeaders: { traceparent?: string; tracestate?: string },
      authorization?: string
    ) => socket.requestWithTracing<Params, Result>(method, params, tracingHeaders, authorization)
  };
}