interface TokenBatch {
  model: string
  provider: string
  count: number
  totalPromptTokens: number
  totalGeneratedTokens: number
  lastEventTime: number
}

export class TokensBatchingService {
  private static instance: TokensBatchingService
  private batches = new Map<string, TokenBatch>()
  private flushTimer: NodeJS.Timeout | null = null

  private readonly BATCH_SIZE_LIMIT = 25
  private readonly FLUSH_INTERVAL_MS = 10 * 60 * 1000 // 10 minutes

  static getInstance(): TokensBatchingService {
    if (!TokensBatchingService.instance) {
      TokensBatchingService.instance = new TokensBatchingService()
    }
    return TokensBatchingService.instance
  }

  private constructor() {
    this.startFlushTimer()
  }

  addTokens(model: string, provider: string, promptTokens: number, generatedTokens: number): void {
    const key = `${provider}:${model}`

    if (!this.batches.has(key)) {
      this.batches.set(key, {
        model,
        provider,
        count: 0,
        totalPromptTokens: 0,
        totalGeneratedTokens: 0,
        lastEventTime: Date.now(),
      })
    }

    const batch = this.batches.get(key)!
    batch.count++
    batch.totalPromptTokens += promptTokens
    batch.totalGeneratedTokens += generatedTokens
    batch.lastEventTime = Date.now()

    // Flush if batch is full
    if (batch.count >= this.BATCH_SIZE_LIMIT) {
      this.flushBatch(key, batch)
    }
  }

  private flushBatch(key: string, batch: TokenBatch): void {
    if (batch.count === 0) return
    this.batches.delete(key)
  }

  private startFlushTimer(): void {
    this.flushTimer = setInterval(() => {
      this.flushAllBatches()
    }, this.FLUSH_INTERVAL_MS)
    // Allow the process to exit if this timer is the only thing keeping it alive
    // This prevents test hangs and allows graceful shutdown
    this.flushTimer.unref()
  }

  private flushAllBatches(): void {
    for (const [key, batch] of this.batches.entries()) {
      this.flushBatch(key, batch)
    }
  }

  shutdown(): void {
    if (this.flushTimer) {
      clearInterval(this.flushTimer)
      this.flushTimer = null
    }
    this.flushAllBatches()
  }
}
