import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient } from './types.js';

export interface RotatingModelClientOptions {
  /** Cooldown period applied to a profile after a failed attempt. */
  cooldownMs?: number;
  /** Injectable clock for deterministic tests. */
  now?: () => number;
}

/**
 * Model client wrapper that rotates across equivalent auth profiles (e.g. API keys).
 * Sticky-by-success behavior: keep using the last successful profile until it fails.
 */
export class RotatingModelClient implements ModelClient {
  private readonly clients: ModelClient[];
  private readonly cooldownMs: number;
  private readonly now: () => number;
  private readonly cooldownUntilMs: number[];
  private currentIndex = 0;

  constructor(clients: ModelClient[], options: RotatingModelClientOptions = {}) {
    if (clients.length === 0) {
      throw new Error('RotatingModelClient requires at least one client');
    }
    this.clients = clients;
    this.cooldownMs = Math.max(0, options.cooldownMs ?? 0);
    this.now = options.now ?? (() => Date.now());
    this.cooldownUntilMs = Array.from({ length: clients.length }, () => 0);
  }

  async chat(request: ChatRequest): Promise<ChatResponse> {
    const errors: Error[] = [];
    const order = this.getAttemptOrder();
    for (const index of order) {
      const client = this.clients[index];
      try {
        const response = await client.chat(request);
        this.currentIndex = index;
        this.cooldownUntilMs[index] = 0;
        return response;
      } catch (error) {
        errors.push(error instanceof Error ? error : new Error(String(error)));
        this.applyCooldown(index);
      }
    }

    throw new Error(`All auth profiles failed: ${errors.map((e) => e.message).join(', ')}`);
  }

  async *chatStream(request: ChatRequest): AsyncIterable<ChatStreamEvent> {
    const order = this.getAttemptOrder();
    for (const index of order) {
      const client = this.clients[index];
      if (!client.chatStream) {
        continue;
      }

      let failed = false;
      for await (const event of client.chatStream(request)) {
        if (event.type === 'error') {
          failed = true;
          this.applyCooldown(index);
          break;
        }
        yield event;
      }

      if (!failed) {
        this.currentIndex = index;
        this.cooldownUntilMs[index] = 0;
        return;
      }
    }

    yield { type: 'error', error: new Error('All auth profiles failed for streaming') };
  }

  private getAttemptOrder(): number[] {
    const now = this.now();
    const available: number[] = [];
    const cooling: number[] = [];

    for (let offset = 0; offset < this.clients.length; offset += 1) {
      const index = (this.currentIndex + offset) % this.clients.length;
      if (this.cooldownUntilMs[index] <= now) {
        available.push(index);
      } else {
        cooling.push(index);
      }
    }

    // If all profiles are cooling down, still attempt them in sticky order.
    return available.length > 0 ? [...available, ...cooling] : cooling;
  }

  private applyCooldown(index: number): void {
    if (this.cooldownMs <= 0) {
      return;
    }
    this.cooldownUntilMs[index] = this.now() + this.cooldownMs;
  }
}