import { InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base";
import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
import { afterEach, beforeEach, describe, expect, it } from "vitest";

import { ClaudeAgentSDKInstrumentation, _resetPatchState } from "../src/instrumentation";

describe("ClaudeAgentSDKInstrumentation", () => {
  let instrumentation: ClaudeAgentSDKInstrumentation;
  let provider: NodeTracerProvider;
  let exporter: InMemorySpanExporter;

  beforeEach(() => {
    exporter = new InMemorySpanExporter();
    provider = new NodeTracerProvider({
      spanProcessors: [new SimpleSpanProcessor(exporter)],
    });
    instrumentation = new ClaudeAgentSDKInstrumentation({
      tracerProvider: provider,
    });
  });

  afterEach(() => {
    if (instrumentation.isEnabled()) {
      instrumentation.disable();
    }
    _resetPatchState();
    exporter.reset();
  });

  it("should be enabled by default", () => {
    expect(instrumentation.isEnabled()).toBe(true);
  });

  it("should be able to enable and disable", () => {
    expect(instrumentation.isEnabled()).toBe(true);

    instrumentation.disable();
    expect(instrumentation.isEnabled()).toBe(false);

    instrumentation.enable();
    expect(instrumentation.isEnabled()).toBe(true);
  });

  it("should initialize without errors", () => {
    expect(() => {
      new ClaudeAgentSDKInstrumentation();
    }).not.toThrow();
  });

  it("should accept custom tracer provider", () => {
    const customProvider = new NodeTracerProvider();
    const customInstrumentation = new ClaudeAgentSDKInstrumentation({
      tracerProvider: customProvider,
    });

    expect(customInstrumentation.isEnabled()).toBe(true);
    customInstrumentation.disable();
  });

  it("should accept trace config options", () => {
    const customInstrumentation = new ClaudeAgentSDKInstrumentation({
      traceConfig: {
        hideInputs: true,
        hideOutputs: true,
      },
    });

    expect(customInstrumentation.isEnabled()).toBe(true);
    customInstrumentation.disable();
  });

  it("should manually instrument a mock module", () => {
    const mockModule = {
      query: function* () {
        yield { type: "result", subtype: "success" };
      },
    };

    expect(() => {
      instrumentation.manuallyInstrument(mockModule);
    }).not.toThrow();
  });

  it("returns a patched copy for native ESM namespace modules", async () => {
    const nativeModule =
      (await import("data:text/javascript,export function query(){return { [Symbol.asyncIterator](){return { async next(){return { done: true } } } } }}")) as unknown as {
        query: () => AsyncIterable<unknown>;
      };

    expect(Object.getOwnPropertyDescriptor(nativeModule, "query")?.writable).toBe(true);
    expect(() => {
      nativeModule.query = nativeModule.query;
    }).toThrow(TypeError);

    const patchedModule = instrumentation.manuallyInstrument(nativeModule);

    expect(patchedModule).not.toBe(nativeModule);
    expect(patchedModule.query).not.toBe(nativeModule.query);
  });

  it("returns a patched copy when SameValue assign works but replacements do not stick", () => {
    const originalQuery = function* () {
      yield { type: "result", subtype: "success" };
    };
    const mockModule = {
      get query() {
        return originalQuery;
      },
      set query(value: typeof originalQuery) {
        // SameValue writes succeed; different values are silently ignored.
        if (value !== originalQuery) {
          return;
        }
      },
    };

    expect(() => {
      mockModule.query = originalQuery;
    }).not.toThrow();
    const otherQuery = function* () {
      yield { type: "result", subtype: "success" };
    };
    mockModule.query = otherQuery;
    expect(mockModule.query).toBe(originalQuery);

    const patchedModule = instrumentation.manuallyInstrument(mockModule);

    expect(patchedModule).not.toBe(mockModule);
    expect(patchedModule.query).not.toBe(originalQuery);
    expect(mockModule.query).toBe(originalQuery);
  });

  it("returns a patched copy when SameValue assign works but replacement throws", () => {
    const originalQuery = function* () {
      yield { type: "result", subtype: "success" };
    };
    const mockModule = {
      get query() {
        return originalQuery;
      },
      set query(value: typeof originalQuery) {
        if (value !== originalQuery) {
          throw new TypeError("Cannot assign to read only property 'query'");
        }
      },
    };

    expect(() => {
      mockModule.query = originalQuery;
    }).not.toThrow();
    expect(() => {
      mockModule.query = function* () {
        yield { type: "result", subtype: "success" };
      };
    }).toThrow(TypeError);

    const patchedModule = instrumentation.manuallyInstrument(mockModule);

    expect(patchedModule).not.toBe(mockModule);
    expect(patchedModule.query).not.toBe(originalQuery);
    expect(mockModule.query).toBe(originalQuery);
  });

  it("should not double-patch the module", () => {
    const callCount = { value: 0 };
    const mockModule = {
      query: function () {
        callCount.value++;
        return {
          [Symbol.asyncIterator]() {
            return {
              async next() {
                return { done: true, value: undefined };
              },
            };
          },
        };
      },
    };

    instrumentation.manuallyInstrument(mockModule);
    const firstQuery = mockModule.query;

    instrumentation.manuallyInstrument(mockModule);
    const secondQuery = mockModule.query;

    // Should not have double-wrapped
    expect(firstQuery).toBe(secondQuery);
  });

  it("should handle V1-only module (no V2 exports)", () => {
    const mockModule = {
      query: function () {
        return {
          [Symbol.asyncIterator]() {
            return {
              async next() {
                return { done: true, value: undefined };
              },
            };
          },
        };
      },
    };

    expect(() => {
      instrumentation.manuallyInstrument(mockModule);
    }).not.toThrow();

    // query should be wrapped
    expect(mockModule.query).not.toBe(undefined);
  });

  it("should handle V2-only module (no query export)", () => {
    const mockModule = {
      unstable_v2_prompt: async (_message: string, _options: Record<string, unknown>) => {
        return {
          type: "result",
          subtype: "success",
          result: "ok",
          usage: { input_tokens: 1, output_tokens: 1 },
          total_cost_usd: 0,
          num_turns: 1,
          duration_ms: 1,
          session_id: "s1",
        };
      },
      unstable_v2_createSession: (_options: Record<string, unknown>) => ({
        sessionId: "s1",
        send: async () => {},
        stream: async function* () {},
        close: () => {},
      }),
    };

    expect(() => {
      instrumentation.manuallyInstrument(mockModule);
    }).not.toThrow();
  });

  it("should handle empty module gracefully", () => {
    const mockModule = {};

    expect(() => {
      instrumentation.manuallyInstrument(mockModule);
    }).not.toThrow();
  });

  it("should not double-wrap after disable() then enable() cycle", () => {
    let _callCount = 0;
    const originalQuery = function () {
      _callCount++;
      return {
        [Symbol.asyncIterator]() {
          return {
            async next() {
              return { done: true, value: undefined };
            },
          };
        },
      };
    };

    const mockModule = {
      query: originalQuery,
    };

    instrumentation.manuallyInstrument(mockModule);
    const wrappedOnce = mockModule.query;
    expect(wrappedOnce).not.toBe(originalQuery);

    // Simulate disable → re-enable cycle
    _resetPatchState();
    // After reset, manuallyInstrument should re-wrap from the original, not double-wrap
    instrumentation.manuallyInstrument(mockModule);
    const wrappedAgain = mockModule.query;

    // Both wrappers should call the original exactly once per invocation
    // We can't easily test single-wrapping without calling, so verify it doesn't throw
    expect(wrappedAgain).toBeDefined();
  });
});
