import type {
  Span,
  SpanContext,
  AttributeValue,
  SpanStatus,
  TimeInput,
  Link,
  Exception,
  Attributes,
  TraceState,
} from "@opentelemetry/api";

export class MockSpan implements Span {
  public attributes: Record<string, AttributeValue> = {};
  public events: {
    name: string;
    attributes?: Attributes | TimeInput;
    startTime?: TimeInput;
  }[] = [];
  public links: Link[] = [];
  public status?: SpanStatus;
  public updatedName?: string;
  public ended: boolean = false;
  public endTime?: TimeInput;
  public exceptions: { exception: Exception; time?: TimeInput }[] = [];

  public resource: { attributes: Attributes } = { attributes: {} };

  spanContext(): SpanContext {
    return {
      traceId: "test-trace-id",
      spanId: "test-span-id",
      traceFlags: 1,
    };
  }

  setAttribute(key: string, value: AttributeValue): this {
    this.attributes[key] = value;
    return this;
  }

  setAttributes(attributes: Attributes): this {
    Object.assign(this.attributes, attributes);
    return this;
  }

  addEvent(
    name: string,
    attributesOrStartTime?: Attributes | TimeInput,
    startTime?: TimeInput,
  ): this {
    this.events.push({ name, attributes: attributesOrStartTime, startTime });
    return this;
  }

  addLink(link: Link): this {
    this.links.push(link);
    return this;
  }

  addLinks(links: Link[]): this {
    this.links.push(...links);
    return this;
  }

  setStatus(status: SpanStatus): this {
    this.status = status;
    return this;
  }

  updateName(name: string): this {
    this.updatedName = name;
    return this;
  }

  end(endTime?: TimeInput): void {
    this.ended = true;
    this.endTime = endTime;
  }

  isRecording(): boolean {
    return true;
  }

  recordException(exception: Exception, time?: TimeInput): void {
    this.exceptions.push({ exception, time });
  }

  reset(): void {
    this.attributes = {};
    this.events = [];
    this.links = [];
    this.status = undefined;
    this.updatedName = undefined;
    this.ended = false;
    this.endTime = undefined;
    this.exceptions = [];
    this.resource = { attributes: {} };
  }
}

export class TestTraceState implements TraceState {
  private entries: Map<string, string>;

  constructor(initEntries?: [string, string][]) {
    this.entries = new Map(initEntries ?? []);
  }

  set(key: string, value: string): TraceState {
    const newEntries: [string, string][] = [[key, value]];
    for (const [k, v] of this.entries.entries()) {
      if (k !== key) {
        newEntries.push([k, v]);
      }
    }
    return new TestTraceState(newEntries);
  }

  unset(key: string): TraceState {
    const newEntries: [string, string][] = [];
    for (const [k, v] of this.entries.entries()) {
      if (k !== key) {
        newEntries.push([k, v]);
      }
    }
    return new TestTraceState(newEntries);
  }

  get(key: string): string | undefined {
    return this.entries.get(key);
  }

  serialize(): string {
    return Array.from(this.entries)
      .slice(0, 32) // Enforce 32-member max
      .map(([key, value]) => `${key}=${value}`)
      .join(",");
  }

  // Optional: for testing/debugging
  toJSON(): Record<string, string> {
    return Object.fromEntries(this.entries);
  }
}