import type { AIUsageSummary, AICallRecord, AICallContext, BaseAICall, LLMUsageData } from './types'
import { calculateCost, LLMModel } from '../../execution/engine/llm/model-info'

/**
 * AIUsageCollector
 * Centralized token tracking that aggregates usage across all LLM calls in an execution
 */
export class AIUsageCollector {
  private model: LLMModel = 'gpt-5' // Default, will be overwritten on first record()
  private calls: AICallRecord[] = []
  private callSequence: number = 0

  /**
   * Record a single AI call with usage metrics
   *
   * @param usage - Token usage and latency data from LLM adapter
   * @param callType - Type discriminator (agent-reasoning, tool, etc.)
   * @param context - Optional typed context specific to callType
   */
  record(usage: LLMUsageData, callType: BaseAICall['callType'] = 'other', context?: AICallContext): void {
    this.callSequence++
    this.model = usage.model

    // Use actual cost from provider (e.g., OpenRouter) if available, otherwise calculate
    const costUsd = usage.cost ?? calculateCost(usage.model, usage.inputTokens, usage.outputTokens)

    this.calls.push({
      callSequence: this.callSequence,
      callType,
      model: usage.model,
      inputTokens: usage.inputTokens,
      outputTokens: usage.outputTokens,
      costUsd,
      latencyMs: usage.latencyMs,
      context
    })
  }

  /**
   * Get aggregated summary of all AI calls
   */
  getSummary(): AIUsageSummary {
    const totalInputTokens = this.calls.reduce((sum, c) => sum + c.inputTokens, 0)
    const totalOutputTokens = this.calls.reduce((sum, c) => sum + c.outputTokens, 0)
    const totalCostUsd = this.calls.reduce((sum, c) => sum + c.costUsd, 0)

    return {
      model: this.model,
      totalInputTokens,
      totalOutputTokens,
      totalTokens: totalInputTokens + totalOutputTokens,
      totalCostUsd,
      callCount: this.calls.length,
      calls: this.calls
    }
  }

  /**
   * Check if any usage has been recorded
   */
  hasUsage(): boolean {
    return this.calls.length > 0
  }
}
