import { Span, SpanOptions, SpanStatusCode, trace } from "@opentelemetry/api";
import { TraceState } from "@opentelemetry/core";
import { NodeSDK } from "@opentelemetry/sdk-node";
import {
  InMemorySpanExporter,
  SimpleSpanProcessor,
  SpanProcessor,
} from "@opentelemetry/sdk-trace-base";
import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest";
import { setNodeSdkConfig } from "../../lib/config-manager.js";
import { getActiveSpan, withSpan } from "../../lib/traces.js";

describe("withSpan", () => {
  let memoryExporter: InMemorySpanExporter;
  let spanProcessor: SpanProcessor;
  let sdk: NodeSDK;

  beforeAll(() => {
    memoryExporter = new InMemorySpanExporter();
    spanProcessor = new SimpleSpanProcessor(memoryExporter);
    sdk = new NodeSDK({
      spanProcessors: [spanProcessor],
      instrumentations: [],
    });

    setNodeSdkConfig({
      collectorUrl: "http://localhost:4317",
    });
    sdk.start();
  });

  afterEach(async () => {
    // Flush any remaining spans
    await spanProcessor.forceFlush();
    // Clean up
    memoryExporter.reset();
  });

  afterAll(async () => {
    await sdk.shutdown();
  });

  it("should handle simple synchronous usage", async ({}) => {
    let capturedSpan: Span;

    await withSpan({
      spanName: "test-sync-span",
      fn: (span: Span) => {
        capturedSpan = span;
      },
    });

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].name).toBe("test-sync-span");
    expect(spans[0].status.code).toBe(SpanStatusCode.OK);
    expect(capturedSpan).toBeTruthy();
    expect(capturedSpan.spanContext().traceId).toBe(
      spans[0].spanContext().traceId,
    );
  });

  it("should handle synchronous functions that throw errors", async ({}) => {
    const error = new Error("Sync error");

    await expect(
      withSpan({
        spanName: "test-sync-error-span",
        fn: () => {
          throw error;
        },
      }),
    ).rejects.toThrow(error.message);

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].name).toBe("test-sync-error-span");
    expect(spans[0].status.code).toBe(SpanStatusCode.ERROR);
    expect(spans[0].status.message).toBe(error.message);
    expect(spans[0].events).toHaveLength(1);
    expect(spans[0].events[0].name).toBe("exception");
  });

  it("should handle asynchronous functions correctly", async ({}) => {
    let capturedSpan: Span;

    await withSpan({
      spanName: "test-async-span",
      fn: async (span: Span) => {
        capturedSpan = span;
        await new Promise((resolve) => setTimeout(resolve, 10));
        return "async-result";
      },
    });

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].name).toBe("test-async-span");
    expect(spans[0].status.code).toBe(SpanStatusCode.OK);
    expect(capturedSpan).toBeTruthy();
    expect(capturedSpan.spanContext().traceId).toBe(
      spans[0].spanContext().traceId,
    );
  });

  it("should handle asynchronous functions that reject", async () => {
    const error = new Error("Async error");

    await expect(
      withSpan({
        spanName: "test-async-error-span",
        fn: async () => {
          await new Promise((resolve) => setTimeout(resolve, 10));
          throw error;
        },
      }),
    ).rejects.toThrow(error.message);

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].name).toBe("test-async-error-span");
    expect(spans[0].status.code).toBe(SpanStatusCode.ERROR);
    expect(spans[0].status.message).toBe(error.message);
    expect(spans[0].events).toHaveLength(1);
    expect(spans[0].events[0].name).toBe("exception");
  });

  it("should handle non-Error exceptions", async () => {
    const nonErrorException = { message: "Not an error object", code: 500 };

    await expect(
      withSpan({
        spanName: "test-non-error-span",
        fn: () => {
          throw nonErrorException;
        },
      }),
    ).rejects.toEqual(nonErrorException);

    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].status.code).toBe(SpanStatusCode.ERROR);
    expect(spans[0].status.message).toBe(JSON.stringify(nonErrorException));
    expect(spans[0].events).toHaveLength(1);
    expect(spans[0].events[0].name).toBe("exception");
  });

  it("should ensure span is ended even when errors occur", async () => {
    const error = new Error("Test error");

    await expect(
      withSpan({
        spanName: "test-finally-span",
        fn: () => {
          throw error;
        },
      }),
    ).rejects.toThrow("Test error");

    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].ended).toBe(true);
  });

  it("should pass span options to tracer", async () => {
    const spanOptions: SpanOptions = {
      attributes: { "custom.attribute": "custom-value" },
      kind: 1,
    };

    await withSpan({
      spanName: "test-options-span",
      spanOptions,
      fn: () => "result",
    });

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].attributes["custom.attribute"]).toBe("custom-value");
    expect(spans[0].kind).toBe(1);
  });

  it("should use custom tracer name", async () => {
    const customTracerName = "custom-tracer";

    await withSpan({
      traceName: customTracerName,
      spanName: "test-custom-tracer-span",
      fn: () => "result",
    });

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].name).toBe("test-custom-tracer-span");
    expect(spans[0].instrumentationScope.name).toBe(customTracerName);
  });

  it("should use default tracer name when not specified", async () => {
    await withSpan({
      spanName: "test-default-tracer-span",
      fn: () => "result",
    });

    await spanProcessor.forceFlush();
    const defaultSpans = memoryExporter.getFinishedSpans();
    expect(defaultSpans).toHaveLength(1);
    expect(defaultSpans[0].name).toBe("test-default-tracer-span");
    expect(defaultSpans[0].instrumentationScope.name).toBe("o11y-sdk");

    memoryExporter.reset();

    setNodeSdkConfig({
      collectorUrl: "",
      serviceName: "test-service",
      serviceVersion: "v1.0.0",
    });

    await withSpan({
      spanName: "test-default-tracer-span",
      fn: () => "result",
    });
    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(spans[0].name).toBe("test-default-tracer-span");
    expect(spans[0].instrumentationScope.name).toBe("test-service");
    expect(spans[0].instrumentationScope.version).toBe("v1.0.0");
  });

  it("should allow function to interact with span context", async () => {
    let receivedSpan: Span;

    await withSpan({
      spanName: "test-span-context",
      fn: (span: Span) => {
        receivedSpan = span;
        span.spanContext().traceState = new TraceState(
          "alpha=aaaaaaaaaaaa,bravo=bbbbbbbbbbbb",
        );
      },
    });

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(1);
    expect(receivedSpan).toBeTruthy();
    expect(receivedSpan.spanContext().traceId).toBe(
      spans[0].spanContext().traceId,
    );
    expect(receivedSpan.spanContext().spanId).toBe(
      spans[0].spanContext().spanId,
    );
    expect(spans[0].spanContext().traceState.serialize()).toStrictEqual(
      "alpha=aaaaaaaaaaaa,bravo=bbbbbbbbbbbb",
    );
  });

  it("should preserve context across setTimeout", async () => {
    await withSpan({
      spanName: "test-timeout-context",
      fn: async (span: Span) => {
        return new Promise((resolve) => {
          setTimeout(() => {
            getActiveSpan().addEvent("promise-resolved", {
              result: "timeout-result",
            });
            resolve("timeout-result");
          }, 10);
        });
      },
    });

    let newSpan: Span;
    await trace
      .getTracer("some-tracer")
      .startActiveSpan("other-context", async (span) => {
        newSpan = span;
        span.addEvent("other-context-event", {
          result: "other-context-result",
        });
      });

    newSpan.addEvent("another-context-event", {
      result: "another-context-result",
    });
    newSpan.setStatus({ code: SpanStatusCode.OK });
    newSpan.end();
    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(2);

    const timeoutSpan = spans.find((s) => s.name === "test-timeout-context");
    expect(timeoutSpan.status.code).toBe(SpanStatusCode.OK);
    expect(timeoutSpan.events).toHaveLength(1);
    expect(timeoutSpan.events[0].name).toStrictEqual("promise-resolved");

    const otherTrackedSpan = spans.find((s) => s.name === "other-context");
    expect(otherTrackedSpan.status.code).toBe(SpanStatusCode.OK);
    expect(otherTrackedSpan.events).toHaveLength(2);
    expect(otherTrackedSpan.events[0].name).toStrictEqual(
      "other-context-event",
    );
    expect(otherTrackedSpan.events[1].name).toStrictEqual(
      "another-context-event",
    );
  });

  it("should handle nested spans correctly", async () => {
    await withSpan({
      spanName: "outer-span",
      fn: async () => {
        await withSpan({
          spanName: "inner-span",
          fn: async () => {
            await new Promise((resolve) => setTimeout(resolve, 1));
          },
        });
      },
    });

    await spanProcessor.forceFlush();
    const spans = memoryExporter.getFinishedSpans();
    expect(spans).toHaveLength(2);

    const innerSpan = spans.find((s) => s.name === "inner-span");
    const outerSpan = spans.find((s) => s.name === "outer-span");

    expect(innerSpan).toBeTruthy();
    expect(outerSpan).toBeTruthy();
    expect(innerSpan!.parentSpanContext.spanId).toBe(
      outerSpan!.spanContext().spanId,
    );
  });
});