app/utils/api_key_pool.py

# ============================================================================
#  API KEY POOL MANAGER
# ============================================================================

import json
import random
import hashlib
import asyncio
from datetime import datetime, timedelta
from typing import Dict, Any, List, Optional
from pathlib import Path

from app.config import Config
from app.utils.logger import logger


class KeyStats:
    """Statistics for a single API key"""
    def __init__(self):
        self.rpm_used = 0
        self.rph_used = 0
        self.rpd_used = 0
        self.tpm_used = 0
        self.tpd_used = 0
        self.total_requests = 0
        self.total_tokens = 0
        self.last_used = None
        self.rpm_reset_at = datetime.now() + timedelta(minutes=1)
        self.rph_reset_at = datetime.now() + timedelta(hours=1)
        self.rpd_reset_at = datetime.now() + timedelta(days=1)
        self.tpm_reset_at = datetime.now() + timedelta(minutes=1)
        self.tpd_reset_at = datetime.now() + timedelta(days=1)


class APIKeyPool:
    """Manages multiple Groq API keys with shuffle rotation and usage tracking"""
    
    USAGE_FILE = Path(Config.OUTPUT_DIR) / "api_usage.json"
    
    def __init__(self, keys: List[str]):
        if not keys:
            raise ValueError("At least one API key is required")
        
        self.keys = keys
        self.key_aliases = {self._hash_key(key): f"key_{i+1}" for i, key in enumerate(keys)}
        self.key_stats: Dict[str, KeyStats] = {key: KeyStats() for key in keys}
        self._lock = asyncio.Lock()
        
        # Load existing usage data
        self._load_usage_data()
        
        logger.info(f"🔐 API Key Pool initialized with {len(keys)} key(s)")
    
    def _hash_key(self, key: str) -> str:
        """Create SHA256 hash of API key for secure storage"""
        return hashlib.sha256(key.encode()).hexdigest()[:16]
    
    def _load_usage_data(self):
        """Load usage data from JSON file"""
        if self.USAGE_FILE.exists():
            try:
                with open(self.USAGE_FILE, 'r') as f:
                    data = json.load(f)
                
                for key in self.keys:
                    key_hash = self._hash_key(key)
                    if key_hash in data.get("keys", {}):
                        stored = data["keys"][key_hash]
                        stats = self.key_stats[key]
                        
                        # Load counters
                        stats.total_requests = stored.get("total_requests", 0)
                        stats.total_tokens = stored.get("total_tokens", 0)
                        stats.last_used = stored.get("last_used")
                        
                        # Load rate limit usage
                        limits = stored.get("limits", {})
                        stats.rpm_used = limits.get("rpm", {}).get("used", 0)
                        stats.rph_used = limits.get("rph", {}).get("used", 0)
                        stats.rpd_used = limits.get("rpd", {}).get("used", 0)
                        stats.tpm_used = limits.get("tpm", {}).get("used", 0)
                        stats.tpd_used = limits.get("tpd", {}).get("used", 0)
                
                logger.info(f"📊 Loaded usage data from {self.USAGE_FILE}")
            except Exception as e:
                logger.warning(f"Failed to load usage data: {e}")
    
    def _save_usage_data(self):
        """Save usage data to JSON file"""
        try:
            data = {
                "last_updated": datetime.now().isoformat(),
                "keys": {}
            }
            
            for key, stats in self.key_stats.items():
                key_hash = self._hash_key(key)
                alias = self.key_aliases[key_hash]
                
                data["keys"][key_hash] = {
                    "alias": alias,
                    "limits": {
                        "rpm": {"used": stats.rpm_used, "limit": Config.GROQ_RPM_LIMIT, "reset_at": stats.rpm_reset_at.isoformat()},
                        "rph": {"used": stats.rph_used, "limit": Config.GROQ_RPH_LIMIT, "reset_at": stats.rph_reset_at.isoformat()},
                        "rpd": {"used": stats.rpd_used, "limit": Config.GROQ_RPD_LIMIT, "reset_at": stats.rpd_reset_at.isoformat()},
                        "tpm": {"used": stats.tpm_used, "limit": Config.GROQ_TPM_LIMIT, "reset_at": stats.tpm_reset_at.isoformat()},
                        "tpd": {"used": stats.tpd_used, "limit": Config.GROQ_TPD_LIMIT, "reset_at": stats.tpd_reset_at.isoformat()}
                    },
                    "stats": {
                        "total_requests": stats.total_requests,
                        "total_tokens": stats.total_tokens,
                        "last_used": stats.last_used,
                        "is_available": self._is_key_available(key)
                    }
                }
            
            # Ensure directory exists
            self.USAGE_FILE.parent.mkdir(parents=True, exist_ok=True)
            
            with open(self.USAGE_FILE, 'w') as f:
                json.dump(data, f, indent=2, default=str)
            
        except Exception as e:
            logger.error(f"Failed to save usage data: {e}")
    
    def _is_key_available(self, key: str) -> bool:
        """Check if key has remaining quota for all limit types"""
        stats = self.key_stats[key]
        now = datetime.now()
        
        # Reset counters if window expired
        if now >= stats.rpm_reset_at:
            stats.rpm_used = 0
            stats.rpm_reset_at = now + timedelta(minutes=1)
        
        if now >= stats.rph_reset_at:
            stats.rph_used = 0
            stats.rph_reset_at = now + timedelta(hours=1)
        
        if now >= stats.rpd_reset_at:
            stats.rpd_used = 0
            stats.rpd_reset_at = now + timedelta(days=1)
        
        if now >= stats.tpm_reset_at:
            stats.tpm_used = 0
            stats.tpm_reset_at = now + timedelta(minutes=1)
        
        if now >= stats.tpd_reset_at:
            stats.tpd_used = 0
            stats.tpd_reset_at = now + timedelta(days=1)
        
        # Check all limits
        if stats.rpm_used >= Config.GROQ_RPM_LIMIT:
            return False
        if stats.rph_used >= Config.GROQ_RPH_LIMIT:
            return False
        if stats.rpd_used >= Config.GROQ_RPD_LIMIT:
            return False
        if stats.tpm_used >= Config.GROQ_TPM_LIMIT:
            return False
        if stats.tpd_used >= Config.GROQ_TPD_LIMIT:
            return False
        
        return True
    
    async def get_available_key(self, estimated_tokens: int = 2500) -> str:
        """Get an available API key using shuffle (random selection)
        
        Args:
            estimated_tokens: Estimated tokens needed for the request
            
        Returns:
            Available API key
            
        Raises:
            RuntimeError: If all keys are at rate limit
        """
        async with self._lock:
            # Filter available keys
            available_keys = [
                key for key in self.keys
                if self._is_key_available(key) and 
                   (self.key_stats[key].tpm_used + estimated_tokens) <= Config.GROQ_TPM_LIMIT and
                   (self.key_stats[key].tpd_used + estimated_tokens) <= Config.GROQ_TPD_LIMIT
            ]
            
            if not available_keys:
                # Find earliest reset time
                earliest_reset = min(
                    stats.tpd_reset_at for stats in self.key_stats.values()
                )
                wait_seconds = (earliest_reset - datetime.now()).total_seconds()
                
                logger.error(f"🚫 All API keys at rate limit. Next reset in {wait_seconds:.0f}s")
                raise RuntimeError(f"All API keys at rate limit. Try again in {wait_seconds:.0f} seconds.")
            
            # SHUFFLE - Random selection from available keys
            selected_key = random.choice(available_keys)
            alias = self.key_aliases[self._hash_key(selected_key)]
            
            logger.info(f"🔀 Selected key: {alias} (from {len(available_keys)} available)")
            
            return selected_key
    
    async def record_usage(self, key: str, tokens_used: int, success: bool = True):
        """Record usage after API call
        
        Args:
            key: The API key used
            tokens_used: Total tokens consumed
            success: Whether the API call succeeded
        """
        async with self._lock:
            stats = self.key_stats[key]
            now = datetime.now()
            
            # Increment request counters
            stats.rpm_used += 1
            stats.rph_used += 1
            stats.rpd_used += 1
            
            # Increment token counters
            stats.tpm_used += tokens_used
            stats.tpd_used += tokens_used
            
            # Update totals
            stats.total_requests += 1
            stats.total_tokens += tokens_used
            stats.last_used = now.isoformat()
            
            # Save to JSON
            self._save_usage_data()
            
            alias = self.key_aliases[self._hash_key(key)]
            logger.info(f"📊 Recorded usage for {alias}: {tokens_used} tokens | RPM: {stats.rpm_used}/{Config.GROQ_RPM_LIMIT} | TPM: {stats.tpm_used}/{Config.GROQ_TPM_LIMIT}")
    
    def get_usage_summary(self) -> Dict[str, Any]:
        """Get usage summary for all keys"""
        summary = {
            "total_keys": len(self.keys),
            "available_keys": sum(1 for k in self.keys if self._is_key_available(k)),
            "total_requests_today": sum(s.rpd_used for s in self.key_stats.values()),
            "total_tokens_today": sum(s.tpd_used for s in self.key_stats.values()),
            "keys": {}
        }
        
        for key, stats in self.key_stats.items():
            key_hash = self._hash_key(key)
            alias = self.key_aliases[key_hash]
            
            summary["keys"][alias] = {
                "is_available": self._is_key_available(key),
                "rpm": {"used": stats.rpm_used, "limit": Config.GROQ_RPM_LIMIT},
                "rph": {"used": stats.rph_used, "limit": Config.GROQ_RPH_LIMIT},
                "rpd": {"used": stats.rpd_used, "limit": Config.GROQ_RPD_LIMIT},
                "tpm": {"used": stats.tpm_used, "limit": Config.GROQ_TPM_LIMIT},
                "tpd": {"used": stats.tpd_used, "limit": Config.GROQ_TPD_LIMIT},
                "total_requests": stats.total_requests,
                "total_tokens": stats.total_tokens,
                "last_used": stats.last_used
            }
        
        return summary
    
    def get_key_alias(self, key: str) -> str:
        """Get alias for a key (for logging/response headers)"""
        key_hash = self._hash_key(key)
        return self.key_aliases.get(key_hash, "unknown")


# Global instance
api_key_pool = APIKeyPool(Config.GROQ_API_KEYS) if Config.GROQ_API_KEYS else None