import type * as AgentsModule from "@openai/agents";
import type { Tracer, TracerProvider } from "@opentelemetry/api";
import { diag } from "@opentelemetry/api";
import type {
  InstrumentationConfig,
  InstrumentationModuleDefinition,
} from "@opentelemetry/instrumentation";
import {
  InstrumentationBase,
  InstrumentationNodeModuleDefinition,
} from "@opentelemetry/instrumentation";

import type { TraceConfigOptions } from "@arizeai/openinference-core";
import { OITracer } from "@arizeai/openinference-core";

import { OpenInferenceTracingProcessor } from "./processor";
// oxlint-disable-next-line typescript/prefer-ts-expect-error
// @ts-ignore - generated by prebuild step
import { VERSION } from "./version";

declare const require: ((id: string) => unknown) | undefined;

export { VERSION };

const MODULE_NAME = "@openai/agents";
const INSTRUMENTATION_NAME = "@arizeai/openinference-instrumentation-openai-agents";

export type OpenAIAgentsInstrumentationOptions = {
  /**
   * When true, replace all SDK trace processors with the OpenInference processor.
   * When false, add the OpenInference processor alongside existing processors.
   * Defaults to true.
   */
  exclusiveProcessor?: boolean;
};

type AgentsModuleWithPatchState = typeof AgentsModule & { openInferencePatched?: boolean };
type InstrumentationModuleState = { moduleExports?: unknown };

/**
 * Flag to track whether the agents module has been patched.
 * Used as a fallback when the module object is immutable.
 */
let _isOpenInferencePatched = false;

/**
 * Check if the agents module is currently instrumented.
 */
export function isPatched(): boolean {
  return _isOpenInferencePatched;
}

/**
 * OpenTelemetry-based instrumentation for the
 * [`@openai/agents`](https://www.npmjs.com/package/@openai/agents) SDK.
 *
 * Registers an {@link OpenInferenceTracingProcessor} with the SDK's native
 * tracing system to emit OpenInference-compliant spans without monkey-patching
 * any SDK methods.
 *
 * @example Auto-instrumentation (CommonJS)
 * ```typescript
 * import { NodeSDK } from "@opentelemetry/sdk-node";
 * import { OpenAIAgentsInstrumentation } from "@arizeai/openinference-instrumentation-openai-agents";
 *
 * const sdk = new NodeSDK({
 *   instrumentations: [new OpenAIAgentsInstrumentation()],
 * });
 * sdk.start();
 * ```
 *
 * @example Manual instrumentation (ESM)
 * ```typescript
 * import * as agents from "@openai/agents";
 * import { OpenAIAgentsInstrumentation } from "@arizeai/openinference-instrumentation-openai-agents";
 *
 * const instrumentation = new OpenAIAgentsInstrumentation({ tracerProvider: provider });
 * instrumentation.manuallyInstrument(agents);
 * ```
 */
export class OpenAIAgentsInstrumentation extends InstrumentationBase<typeof AgentsModule> {
  private oiTracer: OITracer;
  private tracerProvider?: TracerProvider;
  private traceConfig?: TraceConfigOptions;
  private maxRootSpansInFlight?: number;
  private processor?: OpenInferenceTracingProcessor;
  private patchedModule?: AgentsModuleWithPatchState;
  private exclusiveProcessor = true;
  private additiveProcessorRegistered = false;

  constructor({
    exclusiveProcessor,
    instrumentationConfig,
    maxRootSpansInFlight,
    traceConfig,
    tracerProvider,
  }: {
    /**
     * When true (the default), replace all SDK trace processors with the
     * OpenInference processor — the SDK's default OpenAI tracing exporter is
     * removed. When false, add the OpenInference processor alongside existing
     * processors. This is the only way to choose the mode when the module is
     * patched automatically (e.g. via `registerInstrumentations` or NodeSDK);
     * it can also be overridden per-call on {@link instrument} and
     * {@link manuallyInstrument}.
     */
    exclusiveProcessor?: boolean;
    /**
     * Standard OTel instrumentation configuration.
     * @see {@link InstrumentationConfig}
     */
    instrumentationConfig?: InstrumentationConfig;
    /**
     * OpenInference trace configuration for masking/redacting sensitive data.
     * @see {@link TraceConfigOptions}
     */
    traceConfig?: TraceConfigOptions;
    /**
     * Maximum number of root trace spans kept in memory while traces are in flight.
     * Defaults to 1000.
     */
    maxRootSpansInFlight?: number;
    /**
     * Optional custom tracer provider. If not provided, the global provider is used.
     * @see {@link TracerProvider}
     */
    tracerProvider?: TracerProvider;
  } = {}) {
    super(INSTRUMENTATION_NAME, VERSION, Object.assign({}, instrumentationConfig));
    this.tracerProvider = tracerProvider;
    this.traceConfig = traceConfig;
    this.maxRootSpansInFlight = maxRootSpansInFlight;
    this.exclusiveProcessor = exclusiveProcessor ?? true;
    this.oiTracer = new OITracer({
      tracer: this.tracerProvider?.getTracer(INSTRUMENTATION_NAME, VERSION) ?? this.tracer,
      traceConfig,
    });
  }

  protected init(): InstrumentationModuleDefinition<typeof AgentsModule> {
    const module = new InstrumentationNodeModuleDefinition<typeof AgentsModule>(
      MODULE_NAME,
      [">=0.1.0"],
      this.patch.bind(this),
      this.unpatch.bind(this),
    );
    return module;
  }

  /**
   * Manually instruments the agents module. Required when the module is not
   * loaded via `require()` (e.g. ESM environments, bundlers, serverless).
   *
   * @param module - The `@openai/agents` module namespace object.
   * @param options - Controls whether SDK trace processors are replaced or added to.
   */
  manuallyInstrument(
    module: typeof AgentsModule,
    options: OpenAIAgentsInstrumentationOptions = {},
  ): void {
    diag.debug(`Manually instrumenting ${MODULE_NAME}`);
    this.exclusiveProcessor = options.exclusiveProcessor ?? this.exclusiveProcessor;
    this.patch(module);
  }

  /**
   * Registers the OpenInference tracing processor with `@openai/agents`.
   *
   * This method works synchronously in CommonJS environments. In ESM-only
   * environments, pass the imported module namespace to {@link manuallyInstrument}.
   */
  instrument(options: OpenAIAgentsInstrumentationOptions = {}): void {
    this.exclusiveProcessor = options.exclusiveProcessor ?? this.exclusiveProcessor;
    this.enable();

    const module = this.getLoadedModule();
    if (module) {
      this.patch(module);
      return;
    }

    diag.warn(
      `${MODULE_NAME} could not be loaded synchronously; use manuallyInstrument(agents) in ESM environments`,
    );
  }

  /**
   * Removes the registered OpenInference tracing processor from the agents SDK.
   *
   * Only acts on a module this instance patched: an instance that never
   * patched (e.g. because another instance got there first) is a no-op.
   */
  uninstrument(): void {
    const module = this.patchedModule;
    this.disable();
    if (module) {
      this.unpatch(module);
    }
  }

  get tracer(): Tracer {
    if (this.tracerProvider) {
      return this.tracerProvider.getTracer(this.instrumentationName, this.instrumentationVersion);
    }
    return super.tracer;
  }

  setTracerProvider(tracerProvider: TracerProvider): void {
    super.setTracerProvider(tracerProvider);
    this.tracerProvider = tracerProvider;
    this.oiTracer = new OITracer({
      tracer: this.tracer,
      traceConfig: this.traceConfig,
    });

    // A processor created before this call still targets the previous
    // provider. If it is already registered with the agents SDK, swap in a
    // replacement built on the new provider so spans don't keep flowing to
    // the old one.
    const oldProcessor = this.processor;
    this.processor = undefined;
    if (!oldProcessor || !this.patchedModule) {
      return;
    }
    oldProcessor.disable();
    const newProcessor = this.getOrCreateProcessor();
    newProcessor.enable();
    if (this.exclusiveProcessor) {
      this.patchedModule.setTraceProcessors([newProcessor]);
    } else {
      // The SDK has no single-processor removal API: the old (now disabled)
      // processor stays registered; add the replacement alongside it.
      this.patchedModule.addTraceProcessor(newProcessor);
    }
  }

  private patch(module: AgentsModuleWithPatchState, _moduleVersion?: string): typeof AgentsModule {
    diag.debug(`Applying patch for ${MODULE_NAME}@${_moduleVersion ?? "unknown"}`);

    if (module?.openInferencePatched || _isOpenInferencePatched) {
      // Another instrumentation instance (or an earlier call on this one) has
      // already registered a processor; registering a second would duplicate
      // spans, so this call is a no-op.
      diag.debug(`${MODULE_NAME} is already instrumented; skipping`);
      return module;
    }

    const processor = this.getOrCreateProcessor();
    processor.enable();
    if (this.exclusiveProcessor) {
      module.setTraceProcessors([processor]);
      this.additiveProcessorRegistered = false;
    } else if (!this.additiveProcessorRegistered) {
      module.addTraceProcessor(processor);
      this.additiveProcessorRegistered = true;
    }

    this.patchedModule = module;
    _isOpenInferencePatched = true;
    try {
      module.openInferencePatched = true;
    } catch (_e) {
      // Module may be frozen in some environments
    }

    return module;
  }

  private unpatch(module: AgentsModuleWithPatchState, _moduleVersion?: string): void {
    diag.debug(`Removing patch for ${MODULE_NAME}@${_moduleVersion ?? "unknown"}`);

    if (this.exclusiveProcessor) {
      module.setTraceProcessors([]);
      this.additiveProcessorRegistered = false;
    } else {
      // The agents SDK has add/set processor APIs, but no single-processor remove.
      // Disable our processor in place so additive uninstrumentation preserves other processors.
      this.processor?.disable();
    }

    if (this.exclusiveProcessor) {
      this.processor = undefined;
    }
    this.patchedModule = undefined;
    _isOpenInferencePatched = false;
    try {
      module.openInferencePatched = false;
    } catch (_e) {
      // Module may be frozen in some environments
    }
  }

  private getOrCreateProcessor(): OpenInferenceTracingProcessor {
    if (!this.processor) {
      this.processor = new OpenInferenceTracingProcessor({
        tracerProvider: this.tracerProvider,
        traceConfig: this.traceConfig,
        maxRootSpansInFlight: this.maxRootSpansInFlight,
      });
    }
    return this.processor;
  }

  private getLoadedModule(): AgentsModuleWithPatchState | undefined {
    const instrumentedModule = this.getInstrumentedModuleExports();
    if (instrumentedModule) {
      return instrumentedModule;
    }

    if (typeof require !== "function") {
      return;
    }

    try {
      return require(MODULE_NAME) as AgentsModuleWithPatchState;
    } catch (error) {
      diag.debug(`Failed to require ${MODULE_NAME}`, error);
    }
  }

  private getInstrumentedModuleExports(): AgentsModuleWithPatchState | undefined {
    // `_modules` is a private field of InstrumentationBase with no public
    // accessor; guard defensively so an upstream rename degrades to the
    // require() fallback instead of throwing.
    try {
      const modules = (this as unknown as { _modules?: InstrumentationModuleState[] })._modules;
      if (!Array.isArray(modules)) {
        return;
      }
      const moduleExports = modules[0]?.moduleExports;
      return moduleExports as AgentsModuleWithPatchState | undefined;
    } catch (error) {
      diag.debug(`Failed to read instrumented module exports for ${MODULE_NAME}`, error);
    }
  }
}
