key-pool.js

import fs from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';
import { logInfo, logDebug, logError } from './logger.js';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

/**
 * API Key Pool Manager
 * 支持多个API Key的负载均衡和故障转移
 */
class KeyPoolManager {
  constructor() {
    this.keys = [];
    this.currentIndex = 0;
    this.strategy = 'round-robin'; // round-robin | weighted | least-used
    this.healthCheckInterval = 300000; // 5分钟
    this.retryAttempts = 3;
    this.statistics = new Map(); // key_id -> { requests: 0, errors: 0, lastUsed: timestamp }
    this.healthCheckTimer = null;
  }

  /**
   * 加载keys配置
   */
  loadKeys() {
    try {
      const keysPath = path.join(__dirname, 'keys.json');

      if (!fs.existsSync(keysPath)) {
        logInfo('No keys.json found, using single key mode');
        return false;
      }

      const keysData = fs.readFileSync(keysPath, 'utf-8');
      const config = JSON.parse(keysData);

      this.strategy = config.strategy || 'round-robin';
      this.healthCheckInterval = config.health_check_interval || 300000;
      this.retryAttempts = config.retry_attempts || 3;

      // 过滤启用的keys
      this.keys = config.keys.filter(key => key.enabled !== false);

      if (this.keys.length === 0) {
        logInfo('No enabled keys found in keys.json, running in standby mode');
        // 不返回false,允许空密钥池启动,等待UI添加
        this.keys = [];
      }

      // 初始化统计信息
      this.keys.forEach(key => {
        if (!this.statistics.has(key.id)) {
          this.statistics.set(key.id, {
            requests: 0,
            errors: 0,
            lastUsed: 0,
            totalTokens: 0,
            healthy: true
          });
        }
      });

      logInfo(`Key pool initialized with ${this.keys.length} keys`);
      if (this.keys.length > 0) {
        logInfo(`Load balancing strategy: ${this.strategy}`);
      }

      // 启动健康检查
      this.startHealthCheck();

      // 🔥 新增: 启动文件监听,实现热重载
      this.watchKeysFile();

      return true;

    } catch (error) {
      logError('Failed to load keys.json', error);
      return false;
    }
  }

  /**
   * 🔥 新增: 监听keys.json文件变化,实现热重载
   */
  watchKeysFile() {
    if (this.fileWatcher) {
      this.fileWatcher.close();
    }

    const keysPath = path.join(__dirname, 'keys.json');

    try {
      this.fileWatcher = fs.watch(keysPath, (eventType, filename) => {
        if (eventType === 'change') {
          logInfo('🔄 Detected keys.json change, reloading configuration...');
          setTimeout(() => {
            this.reloadKeys();
          }, 100); // 延迟100ms,确保文件写入完成
        }
      });

      logInfo('🔍 File watcher enabled for keys.json (hot reload)');
    } catch (error) {
      logError('Failed to setup file watcher', error);
    }
  }

  /**
   * 🔥 新增: 重新加载密钥配置(不重启应用)
   */
  reloadKeys() {
    try {
      const keysPath = path.join(__dirname, 'keys.json');
      const keysData = fs.readFileSync(keysPath, 'utf-8');
      const config = JSON.parse(keysData);

      // 保存旧的keys用于对比
      const oldKeyIds = new Set(this.keys.map(k => k.id));

      // 更新配置
      this.strategy = config.strategy || 'round-robin';
      this.healthCheckInterval = config.health_check_interval || 300000;
      this.retryAttempts = config.retry_attempts || 3;

      // 过滤启用的keys
      const newKeys = config.keys.filter(key => key.enabled !== false);
      const newKeyIds = new Set(newKeys.map(k => k.id));

      // 检测变化
      const addedKeys = newKeys.filter(k => !oldKeyIds.has(k.id));
      const removedKeys = this.keys.filter(k => !newKeyIds.has(k.id));

      // 更新keys列表
      this.keys = newKeys;

      // 为新增的keys初始化统计信息
      addedKeys.forEach(key => {
        if (!this.statistics.has(key.id)) {
          this.statistics.set(key.id, {
            requests: 0,
            errors: 0,
            lastUsed: 0,
            totalTokens: 0,
            healthy: true
          });
          logInfo(`➕ Added new key: ${key.id}`);
        }
      });

      // 清理已删除的keys统计
      removedKeys.forEach(key => {
        this.statistics.delete(key.id);
        logInfo(`➖ Removed key: ${key.id}`);
      });

      logInfo(`✅ Keys reloaded: ${this.keys.length} keys active (${addedKeys.length} added, ${removedKeys.length} removed)`);

    } catch (error) {
      logError('Failed to reload keys.json', error);
    }
  }

  /**
   * 获取下一个可用的API Key
   * 根据配置的策略选择
   */
  getNextKey() {
    if (this.keys.length === 0) {
      throw new Error('No available keys in pool');
    }

    // 过滤健康的keys
    const healthyKeys = this.keys.filter(key => {
      const stats = this.statistics.get(key.id);
      return stats && stats.healthy;
    });

    if (healthyKeys.length === 0) {
      logError('No healthy keys available, attempting to use all keys');
      // 如果没有健康的key，尝试重置所有key的健康状态
      this.keys.forEach(key => {
        const stats = this.statistics.get(key.id);
        if (stats) stats.healthy = true;
      });
      return this.selectKeyByStrategy(this.keys);
    }

    return this.selectKeyByStrategy(healthyKeys);
  }

  /**
   * 根据策略选择key
   */
  selectKeyByStrategy(availableKeys) {
    let selectedKey;

    switch (this.strategy) {
      case 'weighted':
        selectedKey = this.selectWeightedKey(availableKeys);
        break;

      case 'least-used':
        selectedKey = this.selectLeastUsedKey(availableKeys);
        break;

      case 'round-robin':
      default:
        selectedKey = this.selectRoundRobinKey(availableKeys);
        break;
    }

    // 更新统计
    const stats = this.statistics.get(selectedKey.id);
    if (stats) {
      stats.requests++;
      stats.lastUsed = Date.now();
    }

    logDebug(`Selected key: ${selectedKey.id} (strategy: ${this.strategy})`);
    return selectedKey;
  }

  /**
   * 轮询策略
   */
  selectRoundRobinKey(availableKeys) {
    const key = availableKeys[this.currentIndex % availableKeys.length];
    this.currentIndex = (this.currentIndex + 1) % availableKeys.length;
    return key;
  }

  /**
   * 加权策略
   */
  selectWeightedKey(availableKeys) {
    if (availableKeys.length === 0) {
      throw new Error('No available keys for weighted selection');
    }

    const totalWeight = availableKeys.reduce((sum, key) => sum + (key.weight || 1), 0);
    let random = Math.random() * totalWeight;

    for (const key of availableKeys) {
      random -= (key.weight || 1);
      if (random <= 0) {
        return key;
      }
    }

    return availableKeys[0];
  }

  /**
   * 最少使用策略
   */
  selectLeastUsedKey(availableKeys) {
    return availableKeys.reduce((least, key) => {
      const currentStats = this.statistics.get(key.id);
      const leastStats = this.statistics.get(least.id);

      if (!currentStats) return least;
      if (!leastStats) return key;

      return currentStats.requests < leastStats.requests ? key : least;
    });
  }

  /**
   * 记录key使用失败
   */
  recordKeyFailure(keyId, errorInfo) {
    const stats = this.statistics.get(keyId);
    if (stats) {
      stats.errors++;

      // 如果错误率过高，标记为不健康
      if (stats.errors >= 3) {
        stats.healthy = false;
        const errorMsg = errorInfo?.message || String(errorInfo) || 'Unknown error';
        logError(`Key ${keyId} marked as unhealthy after ${stats.errors} errors: ${errorMsg}`);
      }
    }
  }

  /**
   * 记录key使用成功（包含token统计）
   */
  recordKeySuccess(keyId, tokensUsed = 0) {
    const stats = this.statistics.get(keyId);
    if (stats) {
      // 重置错误计数
      stats.errors = 0;
      stats.healthy = true;
      stats.totalTokens += tokensUsed;
    }
  }

  /**
   * 带重试的key获取
   */
  async getKeyWithRetry() {
    let lastError;
    const attemptedKeys = new Set();

    for (let attempt = 0; attempt < this.retryAttempts; attempt++) {
      try {
        const key = this.getNextKey();

        // 避免重复尝试同一个key
        if (attemptedKeys.has(key.id)) {
          continue;
        }
        attemptedKeys.add(key.id);

        return key;

      } catch (error) {
        lastError = error;
        logDebug(`Key selection attempt ${attempt + 1} failed: ${error.message}`);
      }
    }

    throw new Error(`Failed to get key after ${this.retryAttempts} attempts: ${lastError?.message}`);
  }

  /**
   * 启动健康检查
   */
  startHealthCheck() {
    if (this.healthCheckTimer) {
      clearInterval(this.healthCheckTimer);
    }

    this.healthCheckTimer = setInterval(() => {
      this.performHealthCheck();
    }, this.healthCheckInterval);

    logInfo(`Health check started (interval: ${this.healthCheckInterval}ms)`);
  }

  /**
   * 执行健康检查
   */
  performHealthCheck() {
    logDebug('Performing health check on all keys');

    this.keys.forEach(key => {
      const stats = this.statistics.get(key.id);
      if (!stats) return;

      // 如果一个key长时间未使用且有错误记录，重置其健康状态
      const timeSinceLastUse = Date.now() - stats.lastUsed;
      if (timeSinceLastUse > this.healthCheckInterval && stats.errors > 0) {
        logInfo(`Resetting health status for key ${key.id}`);
        stats.errors = 0;
        stats.healthy = true;
      }
    });
  }

  /**
   * 获取统计信息
   */
  getStatistics() {
    const stats = {
      totalKeys: this.keys.length,
      healthyKeys: 0,
      totalRequests: 0,
      totalErrors: 0,
      totalTokens: 0,
      totalQuota: 0,
      strategy: this.strategy,
      keys: []
    };

    this.keys.forEach(key => {
      const keyStat = this.statistics.get(key.id);
      if (!keyStat) return;

      if (keyStat.healthy) stats.healthyKeys++;
      stats.totalRequests += keyStat.requests;
      stats.totalErrors += keyStat.errors;
      stats.totalTokens += keyStat.totalTokens;
      stats.totalQuota += key.quota || 0;

      stats.keys.push({
        id: key.id,
        healthy: keyStat.healthy,
        requests: keyStat.requests,
        errors: keyStat.errors,
        tokens: keyStat.totalTokens,
        quota: key.quota,
        weight: key.weight,
        errorRate: keyStat.requests > 0 ? (keyStat.errors / keyStat.requests * 100).toFixed(2) + '%' : '0%'
      });
    });

    return stats;
  }

  /**
   * 停止健康检查
   */
  stop() {
    if (this.healthCheckTimer) {
      clearInterval(this.healthCheckTimer);
      this.healthCheckTimer = null;
      logInfo('Health check stopped');
    }

    if (this.fileWatcher) {
      this.fileWatcher.close();
      this.fileWatcher = null;
      logInfo('File watcher stopped');
    }
  }
}

// 单例实例
const keyPool = new KeyPoolManager();

export { keyPool };