import type {
  TracePayload,
  SpanPayload,
  DecisionPointPayload,
  EventPayload,
  JsonValue,
  TraceStatus,
  SpanType,
  SpanStatus,
  DecisionType,
  EventType,
} from "./models.js";
import {
  generateId,
  nowISO,
  TraceStatus as TraceStatusValues,
  SpanStatus as SpanStatusValues,
} from "./models.js";
import { _getTransport } from "./_registry.js";

export interface TraceBuilderOptions {
  sessionId?: string;
  tags?: string[];
  metadata?: JsonValue;
}

export interface AddSpanInput {
  id?: string;
  parentSpanId?: string;
  name: string;
  type: SpanType;
  input?: JsonValue;
  output?: JsonValue;
  tokenCount?: number;
  costUsd?: number;
  durationMs?: number;
  status?: SpanStatus;
  statusMessage?: string;
  startedAt?: string;
  endedAt?: string;
  metadata?: JsonValue;
}

export interface AddDecisionInput {
  id?: string;
  type: DecisionType;
  chosen: JsonValue;
  alternatives?: JsonValue[];
  reasoning?: string;
  contextSnapshot?: JsonValue;
  durationMs?: number;
  costUsd?: number;
  parentSpanId?: string;
  timestamp?: string;
}

export interface AddEventInput {
  id?: string;
  spanId?: string;
  type: EventType;
  name: string;
  metadata?: JsonValue;
  timestamp?: string;
}

export interface EndOptions {
  status?: TraceStatus;
  metadata?: JsonValue;
  totalCost?: number;
  totalTokens?: number;
}

export class TraceBuilder {
  private readonly trace: TracePayload;
  private readonly startMs: number;

  constructor(name: string, options?: TraceBuilderOptions) {
    this.startMs = Date.now();
    this.trace = {
      id: generateId(),
      name,
      sessionId: options?.sessionId,
      status: TraceStatusValues.RUNNING,
      tags: options?.tags ?? [],
      metadata: options?.metadata,
      startedAt: nowISO(),
      decisionPoints: [],
      spans: [],
      events: [],
    };
  }

  addSpan(input: AddSpanInput): string {
    const id = input.id ?? generateId();
    const span: SpanPayload = {
      id,
      parentSpanId: input.parentSpanId,
      name: input.name,
      type: input.type,
      input: input.input,
      output: input.output,
      tokenCount: input.tokenCount,
      costUsd: input.costUsd,
      durationMs: input.durationMs,
      status: input.status ?? SpanStatusValues.RUNNING,
      statusMessage: input.statusMessage,
      startedAt: input.startedAt ?? nowISO(),
      endedAt: input.endedAt,
      metadata: input.metadata,
    };
    this.trace.spans.push(span);
    return id;
  }

  addDecision(input: AddDecisionInput): string {
    const id = input.id ?? generateId();
    const decision: DecisionPointPayload = {
      id,
      type: input.type,
      chosen: input.chosen,
      alternatives: input.alternatives ?? [],
      reasoning: input.reasoning,
      contextSnapshot: input.contextSnapshot,
      durationMs: input.durationMs,
      costUsd: input.costUsd,
      parentSpanId: input.parentSpanId,
      timestamp: input.timestamp ?? nowISO(),
    };
    this.trace.decisionPoints.push(decision);
    return id;
  }

  addEvent(input: AddEventInput): string {
    const id = input.id ?? generateId();
    const event: EventPayload = {
      id,
      spanId: input.spanId,
      type: input.type,
      name: input.name,
      metadata: input.metadata,
      timestamp: input.timestamp ?? nowISO(),
    };
    this.trace.events.push(event);
    return id;
  }

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

  setMetadata(metadata: JsonValue): this {
    this.trace.metadata = metadata;
    return this;
  }

  toPayload(): TracePayload {
    return { ...this.trace };
  }

  end(options?: EndOptions): TracePayload {
    const endedAt = nowISO();
    this.trace.endedAt = endedAt;
    this.trace.totalDuration = Date.now() - this.startMs;
    this.trace.status =
      options?.status ?? TraceStatusValues.COMPLETED;

    if (options?.metadata !== undefined) {
      this.trace.metadata = options.metadata;
    }
    if (options?.totalCost !== undefined) {
      this.trace.totalCost = options.totalCost;
    }
    if (options?.totalTokens !== undefined) {
      this.trace.totalTokens = options.totalTokens;
    }

    const transport = _getTransport();
    if (transport) {
      transport.add(this.trace);
    }

    return { ...this.trace };
  }
}
