import type { ITracer, SpanOpts } from "./tracer.interface";
import type { Instrumented } from "../conformance/brands";
import { attachBrand } from "../conformance/brand-runtime";

const PROPAGATED_BRANDS = [
  "__captured",
  "__audited",
  "__analyzed",
  "__consentChecked",
  "__rateLimited",
] as const;

export function withSpan<Args extends unknown[], R, Extra extends object>(
  tracer: ITracer,
  opts: SpanOpts | ((args: Args) => SpanOpts),
  fn: ((...args: Args) => Promise<R>) & Extra,
): Instrumented<((...args: Args) => Promise<R>) & Extra>;
export function withSpan<Args extends unknown[], R>(
  tracer: ITracer,
  opts: SpanOpts | ((args: Args) => SpanOpts),
  fn: (...args: Args) => Promise<R>,
): Instrumented<(...args: Args) => Promise<R>>;
export function withSpan<Args extends unknown[], R>(
  tracer: ITracer,
  opts: SpanOpts | ((args: Args) => SpanOpts),
  fn: (...args: Args) => Promise<R>,
): Instrumented<(...args: Args) => Promise<R>> {
  const wrapped: (...args: Args) => Promise<R> = (...args) => {
    const resolved = typeof opts === "function" ? opts(args) : opts;
    return tracer.startSpan(resolved, () => fn(...args));
  };
  attachBrand(wrapped, "__instrumented");
  // Propagate brands from the inner function (e.g. __captured from withCapture,
  // __audited from withAudit) so the outermost binding carries all brands.
  // withSpan is always outermost — the assertFeatureConformance check reads the
  // container-resolved value (the withSpan result), so brands must be visible here.
  for (const brand of PROPAGATED_BRANDS) {
    if ((fn as unknown as Record<string, unknown>)[brand] === true) {
      attachBrand(wrapped, brand);
    }
  }
  // Cast is the type-level concession — the brand is now also a non-enumerable
  // runtime property attached above by `attachBrand`.
  return wrapped as Instrumented<(...args: Args) => Promise<R>>;
}