import { Span, SpanOptions, SpanStatusCode, trace } from "@opentelemetry/api";
import { getNodeSdkConfig } from "./config-manager.js";

export type WithSpanParams<T> = {
  /**
   * The name of the trace the span should belong to.
   * NOTE: If you want the new span to belong to an already existing trace, you should provide the same tracer name
   */
  traceName?: string;
  spanName: string;
  spanOptions?: SpanOptions;
  /** A function defining the task you want to be wrapped by this span */
  fn: (span: Span) => T | Promise<T>;
};

/**
 * Generates a function wrapping a given Callable `fn` into an error handling block.
 * Setting Span status and recording any caught exception before bubbling it up.
 *
 * Marks the span as ended once the provided callable has ended or an error has been caught.
 *
 * @returns {Promise<T>} where T is the type returned by the Callable.
 * @throws any error thrown by the original Callable `fn` provided.
 */
function selfContainedSpanHandlerGenerator<T>(
  fn: (span: Span) => T | Promise<T>,
): (span: Span) => Promise<T> {
  return async (span: Span) => {
    try {
      const fnResult = await fn(span);
      span.setStatus({ code: SpanStatusCode.OK });
      return fnResult;
    } catch (err) {
      if (err instanceof Error) {
        span.recordException(err);
        span.setStatus({ code: SpanStatusCode.ERROR, message: err.message });
        throw err;
      }

      span.recordException({ message: JSON.stringify(err) });
      span.setStatus({
        code: SpanStatusCode.ERROR,
        message: JSON.stringify(err),
      });
      throw err;
    } finally {
      span.end();
    }
  };
}

/**
 * Gets the currently active OpenTelemetry span.
 *
 * @returns {Span | undefined} The active span with redaction logic applied,
 * or `undefined` if there is no active span in context.
 */
export function getActiveSpan() {
  return trace.getActiveSpan();
}

export function withSpan<T>({
  traceName,
  spanName,
  spanOptions = {},
  fn,
}: WithSpanParams<T>) {
  const sdkConfig = getNodeSdkConfig();
  const tracer = trace.getTracer(
    traceName ?? sdkConfig.serviceName ?? "o11y-sdk",
    sdkConfig.serviceVersion,
  );
  return tracer.startActiveSpan(
    spanName,
    spanOptions,
    selfContainedSpanHandlerGenerator<T>(fn),
  );
}