import { Tracer } from '@opentelemetry/api';
import WebSocket from 'isomorphic-ws';

import { createISocketClient } from './socket.js';
import { TracedSocket } from './tracedSocket.js';
import { GenericMiddleware, ISocketClient, MethodHandlers, RequestContextBase, SocketTimeouts } from './types.js';
import { connectWebSocket } from './util.js';

// a subclass of ISocket that sends an auth token on the first request
// this is useful for client-side sockets that need to authenticate
export class ISocketWithClientAuth<ImplementedMethods, CallableMethods, RequestContext extends RequestContextBase> extends TracedSocket<
  ImplementedMethods,
  CallableMethods,
  RequestContext
> {
  private authorization?: string;
  private authorizationVersion = 0;
  private hasSentAuth = false;

  constructor(
    ws: WebSocket,
    authorization: string | undefined,
    requestHandlers: MethodHandlers<ImplementedMethods, CallableMethods, RequestContext>,
    globalMiddlewares: GenericMiddleware<CallableMethods, RequestContext>[],
    tracer: Tracer,
    options: {
      onClose: (event: WebSocket.CloseEvent) => void;
      timeouts?: SocketTimeouts;
      logger?: { error: (message?: string) => void; info: (message?: string) => void; warn: (message?: string) => void };
      getAbortSignal?: () => AbortSignal | undefined;
    }
  ) {
    super(ws, requestHandlers, globalMiddlewares, tracer, options);
    this.authorization = authorization;
  }

  // override `request` from the base class to send `authorization` when appropriate
  async request<Params, Result>(method: string, params: Params): Promise<Result> {
    // only send `authorization` on the first request
    const authorizationVersion = this.authorizationVersion;
    const authorization = this.hasSentAuth ? undefined : this.authorization;
    const result = await super.request<Params, Result>(method, params, authorization);
    if (this.authorizationVersion === authorizationVersion) {
      this.hasSentAuth = true;
    }
    return result;
  }

  /**
   * Updates the authorization token for this socket.
   * This is useful for long-lived connections where the token may expire and need to be refreshed.
   *
   * @param newAuthorization - The new authorization token (e.g., "Bearer <token>") or undefined to clear it
   */
  setAuthorization(newAuthorization: string | undefined): void {
    this.authorization = newAuthorization;
    this.authorizationVersion++;
    this.hasSentAuth = false;
  }
}

/**
 * Creates a new socket for a set of operations, without any retry behavior
 * If the requests fail, the promise will reject.
 */
export async function connectSingleUseISocket<
  CallableMethods,
  ImplementedMethods,
  RequestContext extends RequestContextBase = RequestContextBase
>(
  wsUrl: string,
  authorization: string | undefined,
  requestHandlers: MethodHandlers<ImplementedMethods, CallableMethods, RequestContext>,
  globalMiddlewares: GenericMiddleware<CallableMethods, RequestContext>[],
  tracer: Tracer,
  options: {
    onClose: (event: WebSocket.CloseEvent) => void;
    timeouts?: SocketTimeouts;
    logger?: { error: (message?: string) => void; info: (message?: string) => void; warn: (message?: string) => void };
  }
): Promise<ISocketClient<CallableMethods>> {
  const ws = await connectWebSocket(wsUrl);
  const isocket = new ISocketWithClientAuth(ws, authorization, requestHandlers, globalMiddlewares, tracer, options);
  return createISocketClient(isocket);
}