import {
  Attributes,
  ROOT_CONTEXT,
  Span,
  SpanKind,
  SpanStatusCode,
  TimeInput,
  context as otelCtx,
  propagation,
  Tracer,
  trace
} from '@opentelemetry/api';
import { SEMATTRS_MESSAGING_DESTINATION_KIND, SEMATTRS_MESSAGING_SYSTEM } from '@opentelemetry/semantic-conventions';
import WebSocket from 'isomorphic-ws';

import {
  OBS_SOCKET_STATUS_CODE,
  OBS_TAG_API_ID,
  OBS_TAG_APPLICATION_ID,
  OBS_TAG_BRANCH,
  OBS_TAG_COMMIT_ID,
  OBS_TAG_HTTP_ROUTE,
  OBS_TAG_HTTP_STATUS_CODE
} from '../observability/index.js';
import { sanitizeError } from '../tracing/errorSanitizer.js';
import { ISocket, SocketError, SocketMessage } from './socket.js';
import { GenericMiddleware, ISocketClient, MethodHandler, MethodHandlers, RequestContextBase, SocketTimeouts } from './types.js';

function isPromise(obj: unknown): obj is Promise<unknown> {
  return obj !== null && typeof obj === 'object' && typeof (obj as Promise<unknown>).then === 'function';
}

// Taken from https://github.com/gadget-inc/opentelemetry-instrumentations/blob/main/packages/opentelemetry-instrumentation-ws/src/index.ts#L21
const endSpan = (traced: () => unknown | Promise<unknown>, span: Span) => {
  try {
    const result = traced();
    if (isPromise(result)) {
      return Promise.resolve(result)
        .catch((err) => {
          setHttpStatusFromError(span, typeof err === 'string' ? new Error(err) : err);
          throw err;
        })
        .finally(() => span.end());
    } else {
      span.end();
      return result;
    }
  } catch (error) {
    setHttpStatusFromError(span, error);
    span.end();
    throw error;
  }
};

function setHttpStatusFromError(span: Span, error: Error): void {
  // Simplified error handling without server-side dependencies
  span.setAttribute(OBS_TAG_HTTP_STATUS_CODE, 500);
  span.recordException(sanitizeError(error));
  span.setStatus({ code: SpanStatusCode.ERROR, message: error.message });
}

/**
 * @deprecated Use `TracedSocket` from `@superblocksteam/telemetry` instead.
 * This class will be removed in a future version.
 *
 * Migration:
 * ```typescript
 * // Before
 * import { TracedSocket } from '@superblocksteam/shared';
 *
 * // After
 * import { TracedSocket, getTelemetryInstance } from '@superblocksteam/telemetry';
 * const { policyEvaluator } = getTelemetryInstance();
 * new TracedSocket(ws, handlers, middlewares, { requestHeaders, policyEvaluator }, options);
 * ```
 */
export class TracedSocket<ImplementedMethods, CallableMethods, RequestContext extends RequestContextBase> extends ISocket<
  ImplementedMethods,
  CallableMethods,
  RequestContext
> {
  /**
   * The TracedSocket class extends the ISocket class to add tracing capabilities for each connection. It manages
   * spans for handler calls to provide detailed tracing information.
   *
   * Unlike Express, the ISocket handler call chain is fully recursive. This recursion makes it challenging to determine
   * when to create and end spans for specific handlers, as handlers may call other handlers before returning a result.
   *
   * An example span for recursive call
   * Main span ----------------------------------------------------------------------------------
   * Handler1     ----------------------------------------------------------------------------
   * Handler2           ----------------------------------------------------------------------
   *
   * To manage this, we keep a reference to the most recently created span and complete that span when we go one step deeper into the recursion.
   * This approach works for now since we don't support much parallelism, but it will need to be updated when we start handling requests concurrently.
   *
   * An example span for recursive call after the above fixes
   * Main span ----------------------------------------------------------------------------------
   * Handler1     -----------------------------
   * Handler2                                  ------------------------------------------------
   *
   */
  private activeSpanByRequestId = new Map<number, Span>();
  private middlewareSpanByReqId = new Map<number, Span>();
  private lastActiveSpan: Span | undefined;
  private tracer: Tracer;

  constructor(
    ws: WebSocket,
    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, options);
    this.tracer = tracer;
  }

  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> {
    let currentContext = otelCtx.active();
    const middleWareSpan = this.middlewareSpanByReqId.get(ctx.requestId);
    const activeSpan = this.activeSpanByRequestId.get(ctx.requestId);
    // We need to end the middleware span since callHandler is called recursively
    // Otherwise spans will be nested which is not representing actual time taken to execute a middleware
    // Since we are ending the span here, it won't capture error state of nested handlers
    if (middleWareSpan) {
      middleWareSpan.end();
      if (activeSpan) {
        currentContext = trace.setSpan(otelCtx.active(), activeSpan);
      }
    }
    let result;
    await otelCtx.with(currentContext, async () => {
      await this.tracer.startActiveSpan(
        `WS HANDLER ${handler.name === '' ? ctx.method : handler.name}`,
        {
          attributes: {
            [SEMATTRS_MESSAGING_SYSTEM]: 'ws',
            [SEMATTRS_MESSAGING_DESTINATION_KIND]: 'websocket'
          },
          kind: SpanKind.SERVER
        },
        async (span: Span) => {
          this.middlewareSpanByReqId.set(ctx.requestId, span);
          result = await endSpan(() => super.callHandler(handler, params, ctx, client, next), span);
          this.middlewareSpanByReqId.delete(ctx.requestId);
        }
      );
    });
    return result as Result;
  }

  protected async handleMessage(message: SocketMessage): Promise<void> {
    if (message.request) {
      const spanName = message.request.method;
      const requestId = message.request.id;
      const payload = message.request.payload;
      await otelCtx.with(
        propagation.extract(ROOT_CONTEXT, message.request),
        async () =>
          await this.tracer.startActiveSpan(
            `WS SERVER ${spanName}`,
            {
              attributes: {
                [SEMATTRS_MESSAGING_SYSTEM]: 'ws',
                [SEMATTRS_MESSAGING_DESTINATION_KIND]: 'websocket',
                [OBS_TAG_HTTP_ROUTE]: spanName,
                [OBS_TAG_APPLICATION_ID]: payload?.['applicationId'],
                [OBS_TAG_API_ID]: payload?.['apiId'],
                [OBS_TAG_BRANCH]: payload?.['branch'] ?? payload?.['branchName'],
                [OBS_TAG_COMMIT_ID]: payload?.['commitId']
              },
              kind: SpanKind.SERVER
            },
            async (span: Span) => {
              this.lastActiveSpan = span;
              this.activeSpanByRequestId.set(requestId, span);
              const result = await endSpan(() => super.handleMessage(message), span);
              this.activeSpanByRequestId.delete(requestId);
              return result;
            }
          )
      );
    } else if (message.response) {
      this.addEvent('ws.received-response', { ['ws.response.id']: message.response.id });
      return await super.handleMessage(message);
    }
  }

  public request<Params, Result>(method: string, params: Params, authorization?: string): Promise<Result> {
    return this.tracer.startActiveSpan(`WS CLIENT ${method}`, (span): Promise<Result> => {
      this.addEvent('ws.send-request', { ['ws.request.method']: method, ['ws.request.id']: this.nxtRequestId });
      const result: Promise<Result> = super.request(method, params, authorization);
      span.end();
      return result;
    });
  }

  protected decorateToSend(message: SocketMessage): SocketMessage {
    // TODO testing this out
    // Use explicit tracing headers if already provided in the request
    if (message.request?.traceparent || message.request?.tracestate) {
      // Headers already provided - no need to inject from context
      return message;
    }

    // Fallback to context.active() injection for backwards compatibility
    // TODO end
    propagation.inject(otelCtx.active(), message.request);
    return message;
  }

  protected respond<Result>(requestId: number, result: Result): void {
    this.addEvent('ws.send-response', { ['ws.response.id']: requestId });
    return super.respond(requestId, result);
  }

  protected respondError(requestId: number, error: SocketError, exception?: Error): void {
    this.addEvent('ws.send-error', { ['ws.response.id']: requestId });
    if (this.lastActiveSpan) {
      setHttpStatusFromError(this.lastActiveSpan, exception ?? new Error(error.message));
      this.lastActiveSpan.setAttribute(OBS_SOCKET_STATUS_CODE, error.code);
    }
    return super.respondError(requestId, error, exception);
  }

  private addEvent(eventName: string, attributes: Attributes | TimeInput): void {
    // Since we do not have parent requestId in all the methods, we are assuming latest created span is the only active span
    // This assumption can be broken if we start sending multiple request through the same connection.
    if (this.lastActiveSpan) {
      this.lastActiveSpan.addEvent(eventName, attributes, new Date());
    }
  }

  public close(reason?: string): void {
    for (const span of this.activeSpanByRequestId.values()) {
      span.end();
    }
    this.activeSpanByRequestId = new Map();
    this.lastActiveSpan = undefined;
    super.close(reason);
  }
}