import { NoAPIKeysError } from "../models/errors";
import { Logger } from "./Logger";

export interface KeyStateOptions {
  key: string;
  enabled?: boolean;
  cooldownUntil?: Date;
  failures?: number;
  successes?: number;
  metadata?: Record<string, unknown>;
}

/**
 * Internal representation of an API key and its operational lifecycle state.
 * Mapped from Python `utils/key_manager.py` KeyState dataclass.
 */
export class KeyState {
  public key: string;
  public enabled: boolean;
  public cooldownUntil?: Date;
  public failures: number;
  public successes: number;
  public metadata: Record<string, unknown>;

  constructor(options: KeyStateOptions) {
    this.key = options.key;
    this.enabled = options.enabled ?? true;
    this.cooldownUntil = options.cooldownUntil;
    this.failures = options.failures ?? 0;
    this.successes = options.successes ?? 0;
    this.metadata = options.metadata ?? {};
  }
}

export interface KeyManagerOptions {
  providerName: string;
  keys?: string[];
  defaultCooldownHours?: number;
}

/**
 * Manager for search provider API keys.
 * Handles round-robin key rotation, key state tracking, HTTP 401 permanent invalidation,
 * and HTTP 403/429 quota exhaustion cooldowns.
 *
 * Key Lifecycle Rules & Configuration Synchronization:
 * 1. Initial State / Plugin Restart:
 *    All configured keys start enabled (`enabled: true`) in memory and are validated on first request.
 * 2. HTTP 401 (Unauthorized / Invalid Key):
 *    Calling `markKeyInvalid()` sets `enabled: false`. The invalid key is skipped on all subsequent
 *    search requests until a configuration update occurs.
 * 3. Configuration Update (`setKeys`):
 *    When settings are edited in LM Studio GUI, `setKeys(newKeys)` replaces in-memory key states with
 *    fresh `KeyState({ key, enabled: true })` instances. This re-enables updated keys for re-validation
 *    on the next search attempt. If the key is still invalid, it will be re-disabled on first use.
 * 4. HTTP 403/429 (Quota / Rate Limit):
 *    Calling `markKeyExhausted()` sets a temporary `cooldownUntil` timestamp (24h or `Retry-After`).
 *    Requests during the cooldown window skip the key without making HTTP calls. When `cooldownUntil`
 *    expires, the key automatically re-enables.
 * 5. Concurrency & In-Flight Requests:
 *    Cooldowns prevent future key selections. Requests already in flight prior to a cooldown being
 *    recorded will complete using their previously selected key without interruption.
 */
export class KeyManager {
  public readonly providerName: string;
  public readonly defaultCooldownHours: number;
  private index: number = 0;
  private keys: KeyState[] = [];

  constructor(options: KeyManagerOptions) {
    this.providerName = options.providerName;
    this.defaultCooldownHours = Math.max(0, options.defaultCooldownHours ?? 24.0);
    if (options.keys && options.keys.length > 0) {
      this.keys = options.keys.map((k) => new KeyState({ key: k, enabled: true }));
    }
  }

  /**
   * Replaces current key snapshot with fresh key strings from configuration update.
   * Resets key states to `enabled: true` for re-validation on next execution.
   */
  public setKeys(keys: string[]): void {
    this.keys = keys.map((k) => new KeyState({ key: k, enabled: true }));
    this.index = 0;
  }

  /**
   * Return the current active API key without rotating.
   * Mapped from Python `KeyManager.current_key()`.
   */
  public currentKey(): string {
    const activeKeys = this.getActiveKeys();
    this.validateHasKeys(activeKeys);
    this.index = this.index % activeKeys.length;
    Logger.debug(
      "[KeyManager]",
      `${this.providerName}: Using key #${this.index + 1} of ${activeKeys.length}`
    );
    return activeKeys[this.index].key;
  }

  /**
   * Rotate to the next available API key in round-robin order and return it.
   * Mapped from Python `KeyManager.next_key()`.
   */
  public nextKey(): string {
    const activeKeys = this.getActiveKeys();
    this.validateHasKeys(activeKeys);
    this.advanceIndex(activeKeys.length);
    Logger.debug(
      "[KeyManager]",
      `${this.providerName}: Rotating to key #${this.index + 1} of ${activeKeys.length}`
    );
    return activeKeys[this.index].key;
  }

  /**
   * Mark an API key as permanently invalid (HTTP 401 Unauthorized / Invalid Key).
   * Key remains disabled until settings update or key replacement.
   */
  public markKeyInvalid(key: string): void {
    let keyIndex = 1;
    for (let i = 0; i < this.keys.length; i++) {
      if (this.keys[i].key === key) {
        keyIndex = i + 1;
        this.keys[i].enabled = false;
        this.keys[i].failures += 1;
        break;
      }
    }
    Logger.error(
      "[KeyManager]",
      `${this.providerName}: Key #${keyIndex} (401 Unauthorized / Invalid API Key) permanently disabled until configuration update.`
    );
  }

  /**
   * Mark an API key as quota exhausted and place it into temporary cooldown.
   * Mapped from Python `KeyManager.mark_key_exhausted()`.
   */
  public markKeyExhausted(key: string, retryAfterSeconds?: number): void {
    let keyIndex = 1;
    let untilDate: Date = new Date();
    for (let i = 0; i < this.keys.length; i++) {
      if (this.keys[i].key === key) {
        keyIndex = i + 1;
        this.keys[i].failures += 1;
        const now = new Date();
        const cooldownMs = (retryAfterSeconds && retryAfterSeconds > 0)
          ? retryAfterSeconds * 1000
          : this.defaultCooldownHours * 3600 * 1000;
        untilDate = new Date(now.getTime() + cooldownMs);
        this.keys[i].cooldownUntil = untilDate;
        break;
      }
    }
    const formattedUntil = untilDate.toISOString().replace("T", " ").substring(0, 16);
    Logger.warn(
      "[KeyManager]",
      `${this.providerName}: Key #${keyIndex} quota exhausted. Disabled until ${formattedUntil}`
    );
  }

  /**
   * Record successful request for the given API key.
   * Mapped from Python `KeyManager.mark_key_success()`.
   */
  public markKeySuccess(key: string): void {
    for (const k of this.keys) {
      if (k.key === key) {
        k.successes += 1;
        break;
      }
    }
  }

  /**
   * Return the total number of configured API keys.
   * Mapped from Python `KeyManager.key_count()`.
   */
  public keyCount(): number {
    return this.keys.length;
  }

  /**
   * Check if any API keys are configured for the provider.
   * Mapped from Python `KeyManager.has_keys()`.
   */
  public hasKeys(): boolean {
    return this.keyCount() > 0;
  }

  /**
   * Reset key rotation back to the first key.
   * Mapped from Python `KeyManager.reset()`.
   */
  public reset(): void {
    this.index = 0;
  }

  /**
   * Internal helper to check if a key state is available.
   * Mapped from Python `KeyManager._is_available()`.
   */
  private isAvailable(keyState: KeyState, keyIndex?: number): boolean {
    if (!keyState.enabled) {
      return false;
    }
    if (keyState.cooldownUntil) {
      const now = new Date();
      if (now < keyState.cooldownUntil) {
        const formattedUntil = keyState.cooldownUntil.toISOString().replace("T", " ").substring(0, 16);
        Logger.debug(
          "[KeyManager]",
          `${this.providerName}: Skipping Key #${keyIndex ?? "?"} (cooldown active until ${formattedUntil})`
        );
        return false;
      }
      keyState.cooldownUntil = undefined;
      Logger.info(
        "[KeyManager]",
        `${this.providerName}: Key #${keyIndex ?? "?"} cooldown expired. Re-enabled for search requests.`
      );
    }
    return true;
  }

  /**
   * Internal helper to retrieve all currently available keys.
   * Mapped from Python `KeyManager._get_active_keys()`.
   */
  private getActiveKeys(): KeyState[] {
    return this.keys.filter((k, idx) => this.isAvailable(k, idx + 1));
  }

  /**
   * Internal helper to step round-robin index pointer.
   * Mapped from Python `KeyManager._advance_index()`.
   */
  private advanceIndex(activeCount: number): void {
    this.index = (this.index + 1) % activeCount;
  }

  /**
   * Internal helper to validate that active keys exist.
   * Mapped from Python `KeyManager._validate_has_keys()`.
   */
  private validateHasKeys(activeKeys: KeyState[]): void {
    if (activeKeys.length === 0) {
      Logger.warn("[Provider]", `No API keys configured or available for ${this.providerName}`);
      throw new NoAPIKeysError(
        `No API keys configured or available for provider '${this.providerName}'`
      );
    }
  }
}

import { NoAPIKeysError } from "../models/errors";
import { Logger } from "./Logger";

export interface KeyStateOptions {
  key: string;
  enabled?: boolean;
  cooldownUntil?: Date;
  failures?: number;
  successes?: number;
  metadata?: Record<string, unknown>;
}

/**
 * Internal representation of an API key and its operational lifecycle state.
 * Mapped from Python `utils/key_manager.py` KeyState dataclass.
 */
export class KeyState {
  public key: string;
  public enabled: boolean;
  public cooldownUntil?: Date;
  public failures: number;
  public successes: number;
  public metadata: Record<string, unknown>;

  constructor(options: KeyStateOptions) {
    this.key = options.key;
    this.enabled = options.enabled ?? true;
    this.cooldownUntil = options.cooldownUntil;
    this.failures = options.failures ?? 0;
    this.successes = options.successes ?? 0;
    this.metadata = options.metadata ?? {};
  }
}

export interface KeyManagerOptions {
  providerName: string;
  keys?: string[];
  defaultCooldownHours?: number;
}

/**
 * Manager for search provider API keys.
 * Handles round-robin key rotation, key state tracking, HTTP 401 permanent invalidation,
 * and HTTP 403/429 quota exhaustion cooldowns.
 *
 * Key Lifecycle Rules & Configuration Synchronization:
 * 1. Initial State / Plugin Restart:
 *    All configured keys start enabled (`enabled: true`) in memory and are validated on first request.
 * 2. HTTP 401 (Unauthorized / Invalid Key):
 *    Calling `markKeyInvalid()` sets `enabled: false`. The invalid key is skipped on all subsequent
 *    search requests until a configuration update occurs.
 * 3. Configuration Update (`setKeys`):
 *    When settings are edited in LM Studio GUI, `setKeys(newKeys)` replaces in-memory key states with
 *    fresh `KeyState({ key, enabled: true })` instances. This re-enables updated keys for re-validation
 *    on the next search attempt. If the key is still invalid, it will be re-disabled on first use.
 * 4. HTTP 403/429 (Quota / Rate Limit):
 *    Calling `markKeyExhausted()` sets a temporary `cooldownUntil` timestamp (24h or `Retry-After`).
 *    Requests during the cooldown window skip the key without making HTTP calls. When `cooldownUntil`
 *    expires, the key automatically re-enables.
 * 5. Concurrency & In-Flight Requests:
 *    Cooldowns prevent future key selections. Requests already in flight prior to a cooldown being
 *    recorded will complete using their previously selected key without interruption.
 */
export class KeyManager {
  public readonly providerName: string;
  public readonly defaultCooldownHours: number;
  private index: number = 0;
  private keys: KeyState[] = [];

  constructor(options: KeyManagerOptions) {
    this.providerName = options.providerName;
    this.defaultCooldownHours = Math.max(0, options.defaultCooldownHours ?? 24.0);
    if (options.keys && options.keys.length > 0) {
      this.keys = options.keys.map((k) => new KeyState({ key: k, enabled: true }));
    }
  }

  /**
   * Replaces current key snapshot with fresh key strings from configuration update.
   * Resets key states to `enabled: true` for re-validation on next execution.
   */
  public setKeys(keys: string[]): void {
    this.keys = keys.map((k) => new KeyState({ key: k, enabled: true }));
    this.index = 0;
  }

  /**
   * Return the current active API key without rotating.
   * Mapped from Python `KeyManager.current_key()`.
   */
  public currentKey(): string {
    const activeKeys = this.getActiveKeys();
    this.validateHasKeys(activeKeys);
    this.index = this.index % activeKeys.length;
    Logger.debug(
      "[KeyManager]",
      `${this.providerName}: Using key #${this.index + 1} of ${activeKeys.length}`
    );
    return activeKeys[this.index].key;
  }

  /**
   * Rotate to the next available API key in round-robin order and return it.
   * Mapped from Python `KeyManager.next_key()`.
   */
  public nextKey(): string {
    const activeKeys = this.getActiveKeys();
    this.validateHasKeys(activeKeys);
    this.advanceIndex(activeKeys.length);
    Logger.debug(
      "[KeyManager]",
      `${this.providerName}: Rotating to key #${this.index + 1} of ${activeKeys.length}`
    );
    return activeKeys[this.index].key;
  }

  /**
   * Mark an API key as permanently invalid (HTTP 401 Unauthorized / Invalid Key).
   * Key remains disabled until settings update or key replacement.
   */
  public markKeyInvalid(key: string): void {
    let keyIndex = 1;
    for (let i = 0; i < this.keys.length; i++) {
      if (this.keys[i].key === key) {
        keyIndex = i + 1;
        this.keys[i].enabled = false;
        this.keys[i].failures += 1;
        break;
      }
    }
    Logger.error(
      "[KeyManager]",
      `${this.providerName}: Key #${keyIndex} (401 Unauthorized / Invalid API Key) permanently disabled until configuration update.`
    );
  }

  /**
   * Mark an API key as quota exhausted and place it into temporary cooldown.
   * Mapped from Python `KeyManager.mark_key_exhausted()`.
   */
  public markKeyExhausted(key: string, retryAfterSeconds?: number): void {
    let keyIndex = 1;
    let untilDate: Date = new Date();
    for (let i = 0; i < this.keys.length; i++) {
      if (this.keys[i].key === key) {
        keyIndex = i + 1;
        this.keys[i].failures += 1;
        const now = new Date();
        const cooldownMs = (retryAfterSeconds && retryAfterSeconds > 0)
          ? retryAfterSeconds * 1000
          : this.defaultCooldownHours * 3600 * 1000;
        untilDate = new Date(now.getTime() + cooldownMs);
        this.keys[i].cooldownUntil = untilDate;
        break;
      }
    }
    const formattedUntil = untilDate.toISOString().replace("T", " ").substring(0, 16);
    Logger.warn(
      "[KeyManager]",
      `${this.providerName}: Key #${keyIndex} quota exhausted. Disabled until ${formattedUntil}`
    );
  }

  /**
   * Record successful request for the given API key.
   * Mapped from Python `KeyManager.mark_key_success()`.
   */
  public markKeySuccess(key: string): void {
    for (const k of this.keys) {
      if (k.key === key) {
        k.successes += 1;
        break;
      }
    }
  }

  /**
   * Return the total number of configured API keys.
   * Mapped from Python `KeyManager.key_count()`.
   */
  public keyCount(): number {
    return this.keys.length;
  }

  /**
   * Check if any API keys are configured for the provider.
   * Mapped from Python `KeyManager.has_keys()`.
   */
  public hasKeys(): boolean {
    return this.keyCount() > 0;
  }

  /**
   * Reset key rotation back to the first key.
   * Mapped from Python `KeyManager.reset()`.
   */
  public reset(): void {
    this.index = 0;
  }

  /**
   * Internal helper to check if a key state is available.
   * Mapped from Python `KeyManager._is_available()`.
   */
  private isAvailable(keyState: KeyState, keyIndex?: number): boolean {
    if (!keyState.enabled) {
      return false;
    }
    if (keyState.cooldownUntil) {
      const now = new Date();
      if (now < keyState.cooldownUntil) {
        const formattedUntil = keyState.cooldownUntil.toISOString().replace("T", " ").substring(0, 16);
        Logger.debug(
          "[KeyManager]",
          `${this.providerName}: Skipping Key #${keyIndex ?? "?"} (cooldown active until ${formattedUntil})`
        );
        return false;
      }
      keyState.cooldownUntil = undefined;
      Logger.info(
        "[KeyManager]",
        `${this.providerName}: Key #${keyIndex ?? "?"} cooldown expired. Re-enabled for search requests.`
      );
    }
    return true;
  }

  /**
   * Internal helper to retrieve all currently available keys.
   * Mapped from Python `KeyManager._get_active_keys()`.
   */
  private getActiveKeys(): KeyState[] {
    return this.keys.filter((k, idx) => this.isAvailable(k, idx + 1));
  }

  /**
   * Internal helper to step round-robin index pointer.
   * Mapped from Python `KeyManager._advance_index()`.
   */
  private advanceIndex(activeCount: number): void {
    this.index = (this.index + 1) % activeCount;
  }

  /**
   * Internal helper to validate that active keys exist.
   * Mapped from Python `KeyManager._validate_has_keys()`.
   */
  private validateHasKeys(activeKeys: KeyState[]): void {
    if (activeKeys.length === 0) {
      Logger.warn("[Provider]", `No API keys configured or available for ${this.providerName}`);
      throw new NoAPIKeysError(
        `No API keys configured or available for provider '${this.providerName}'`
      );
    }
  }
}