import { Context, context, Span, trace, Tracer } from '@opentelemetry/api';

import { sanitizeError } from './errorSanitizer.js';

/**
 * helper function to wrap a class method with tracing
 * @param className - The class name
 * @param methodName - The method name
 * @param descriptor - The method descriptor
 * @param tracer - The tracer
 */
export function traceMethod(className: string, methodName: string, descriptor: PropertyDescriptor, tracer: Tracer): PropertyDescriptor {
  const originalMethod = descriptor.value;

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  descriptor.value = function (...args: any[]) {
    // Potentially, we can pass this from tracing middleware to req and then to ActionHandler
    const parentSpan = trace.getActiveSpan() as Span | undefined;
    let ctx: Context;
    if (parentSpan) {
      ctx = trace.setSpan(context.active(), parentSpan);
    } else {
      ctx = context.active();
    }

    // Use `startActiveSpan` to start and manage the span lifecycle
    return context.with(ctx, () => {
      return tracer.startActiveSpan(`${className}.${methodName}`, {}, async (span) => {
        try {
          const result = originalMethod.apply(this, args);

          if (result instanceof Promise) {
            return result
              .then((res) => {
                span.end();
                return res;
              })
              .catch((err) => {
                span.recordException(sanitizeError(err));
                span.end();
                throw err;
              });
          }

          // For synchronous methods, end the span and return the result
          span.end();
          return result;
        } catch (err) {
          span.recordException(sanitizeError(err));
          span.end();
          throw err;
        }
      });
    });
  };
  return descriptor;
}

export function traceFunction<T>({
  name,
  tracer,
  fn,
  parent
}: {
  name: string;
  tracer: Tracer;
  fn: (span: Span) => T | Promise<T>;
  parent?: Span;
}): Promise<T> {
  const parentSpan = parent ?? (trace.getActiveSpan() as Span | undefined);
  const ctx = parentSpan ? trace.setSpan(context.active(), parentSpan) : context.active();

  return context.with(ctx, () =>
    tracer.startActiveSpan(name, {}, async (span) => {
      try {
        const result = fn(span);

        if (result instanceof Promise) {
          return result
            .then((res) => {
              span.end();
              return res;
            })
            .catch((err) => {
              span.recordException(sanitizeError(err));
              span.end();
              throw err;
            });
        }

        span.end();
        return result;
      } catch (err) {
        span.recordException(sanitizeError(err));
        span.end();
        throw err;
      }
    })
  );
}

/**
 * @deprecated Use `traceFunction` instead.
 */
export function executeTraced<T>(name: string, tracer: Tracer, fn: (span: Span) => Promise<T>, parent?: Span): Promise<T> {
  const parentSpan = parent ?? (trace.getActiveSpan() as Span | undefined);
  if (parentSpan) {
    // use parent span as context
    return context.with(trace.setSpan(context.active(), parentSpan), () => {
      const span = tracer.startSpan(name);
      return fn(span).finally(() => {
        span.end();
      });
    });
  } else {
    // use new span as context
    return context.with(context.active(), () => {
      const span = tracer.startSpan(name);
      return fn(span).finally(() => {
        span.end();
      });
    });
  }
}

function findDescriptor(object: object, method: string): PropertyDescriptor | undefined {
  let currentProto = Object.getPrototypeOf(object);

  // Traverse up the prototype chain to find the method's descriptor
  while (currentProto) {
    const descriptor = Object.getOwnPropertyDescriptor(currentProto, method);
    if (descriptor) {
      return descriptor; // Return the found descriptor
    }
    currentProto = Object.getPrototypeOf(currentProto); // Move up the prototype chain
  }

  return undefined; // Return undefined if not found
}

export function addTracingToMethods(object: object, methods: string[], tracer: Tracer) {
  methods.forEach((method) => {
    if (typeof object[method] === 'function') {
      const descriptor = findDescriptor(object, method);
      if (descriptor) {
        Object.defineProperty(object, method, traceMethod(object.constructor.name, method, descriptor, tracer));
      }
    }
  });
}