api_key_manager.py

#!/usr/bin/env python3
"""
API Key Manager for MangaTranslator
===================================

Quản lý nhiều API key và tự động rotate để tối ưu usage.

Features:
- Hỗ trợ nhiều provider (Gemini, Google Translate, etc.)
- Round-robin rotation hoặc random selection
- Tracking usage và daily limits
- Auto fallback khi key hết quota
- Reset daily usage tự động

Author: MangaTranslator Team
License: MIT
Version: 1.0
"""

import json
import os
import random
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
import time

class ApiKeyManager:
    """
    Quản lý và rotate API keys cho các dịch vụ translation
    """
    
    def __init__(self, config_path: str = "config/api_keys.json"):
        """
        Initialize API Key Manager
        
        Args:
            config_path (str): Đường dẫn đến file config JSON
        """
        self.config_path = config_path
        self.config = {}
        self.current_indices = {}  # Track current index for round-robin
        self.last_used_times = {}  # Track last used time for rate limiting
        
        self.load_config()
        
    def load_config(self):
        """Load configuration từ JSON file"""
        try:
            if os.path.exists(self.config_path):
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    self.config = json.load(f)
                    
                # Reset daily usage nếu cần
                if self.config.get('auto_reset_usage', True):
                    self._reset_daily_usage_if_needed()
                    
                print(f"✅ Loaded API keys config from {self.config_path}")
                self._print_key_status()
            else:
                print(f"⚠️ Config file not found: {self.config_path}")
                self._create_default_config()
                
        except Exception as e:
            print(f"❌ Error loading config: {e}")
            self.config = {}
    
    def save_config(self):
        """Lưu configuration vào JSON file"""
        try:
            os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
            with open(self.config_path, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, indent=2, ensure_ascii=False)
                
        except Exception as e:
            print(f"❌ Error saving config: {e}")
    
    def _create_default_config(self):
        """Tạo config mặc định nếu file không tồn tại"""
        default_config = {
            "gemini_keys": [],
            "google_translate_keys": [],
            "rotation_strategy": "round_robin",
            "fallback_to_free": True,
            "auto_reset_usage": True
        }
        
        self.config = default_config
        self.save_config()
        print(f"📝 Created default config at {self.config_path}")
    
    def _reset_daily_usage_if_needed(self):
        """Reset daily usage nếu đã sang ngày mới"""
        today = datetime.now().strftime("%Y-%m-%d")
        
        for provider in ['gemini_keys', 'google_translate_keys']:
            if provider in self.config:
                for key_info in self.config[provider]:
                    if key_info.get('last_reset') != today:
                        key_info['usage_count'] = 0
                        key_info['last_reset'] = today
        
        # Lưu config sau khi reset
        self.save_config()
    
    def _print_key_status(self):
        """In ra status của các API keys"""
        print("\n📊 API Key Status:")
        
        for provider in ['gemini_keys', 'google_translate_keys']:
            if provider in self.config and self.config[provider]:
                provider_name = provider.replace('_keys', '').title()
                print(f"\n{provider_name}:")
                
                for i, key_info in enumerate(self.config[provider]):
                    status = "🟢" if key_info.get('active', False) else "🔴"
                    usage = key_info.get('usage_count', 0)
                    limit = key_info.get('daily_limit', 0)
                    name = key_info.get('name', f'Key {i+1}')
                    
                    print(f"  {status} {name}: {usage}/{limit} requests")
    
    def get_api_key(self, provider: str) -> Optional[str]:
        """
        Lấy API key cho provider theo rotation strategy
        
        Args:
            provider (str): Tên provider ('gemini', 'google_translate')
            
        Returns:
            str: API key hoặc None nếu không có key khả dụng
        """
        provider_key = f"{provider}_keys"
        
        if provider_key not in self.config or not self.config[provider_key]:
            print(f"⚠️ No {provider} keys configured")
            return None
        
        keys = self.config[provider_key]
        available_keys = [k for k in keys if k.get('active', False) and not self._is_key_exhausted(k)]
        
        if not available_keys:
            print(f"⚠️ No available {provider} keys (all exhausted or inactive)")
            return None
        
        # Chọn key theo strategy
        strategy = self.config.get('rotation_strategy', 'round_robin')
        
        if strategy == 'round_robin':
            selected_key = self._get_key_round_robin(provider, available_keys)
        elif strategy == 'random':
            selected_key = random.choice(available_keys)
        elif strategy == 'least_used':
            selected_key = min(available_keys, key=lambda k: k.get('usage_count', 0))
        else:
            selected_key = available_keys[0]  # Default to first
        
        if selected_key:
            # Increment usage count
            selected_key['usage_count'] = selected_key.get('usage_count', 0) + 1
            self.save_config()
            
            key_name = selected_key.get('name', 'Unknown')
            usage = selected_key.get('usage_count', 0)
            limit = selected_key.get('daily_limit', 0)
            
            print(f"🔑 Using {provider} key: {key_name} ({usage}/{limit})")
            
            return selected_key['key']
        
        return None
    
    def _get_key_round_robin(self, provider: str, available_keys: List[Dict]) -> Optional[Dict]:
        """Get key using round-robin strategy"""
        if provider not in self.current_indices:
            self.current_indices[provider] = 0
        
        # Find the original index in the config
        all_keys = self.config[f"{provider}_keys"]
        current_key = available_keys[self.current_indices[provider] % len(available_keys)]
        
        # Update index for next call
        self.current_indices[provider] = (self.current_indices[provider] + 1) % len(available_keys)
        
        return current_key
    
    def _is_key_exhausted(self, key_info: Dict) -> bool:
        """Kiểm tra xem key đã hết quota chưa"""
        usage = key_info.get('usage_count', 0)
        limit = key_info.get('daily_limit', float('inf'))
        
        return usage >= limit
    
    def add_api_key(self, provider: str, key: str, name: str = None, daily_limit: int = 1000):
        """
        Thêm API key mới
        
        Args:
            provider (str): Provider name ('gemini', 'google_translate')
            key (str): API key
            name (str): Tên mô tả cho key
            daily_limit (int): Giới hạn sử dụng hàng ngày
        """
        provider_key = f"{provider}_keys"
        
        if provider_key not in self.config:
            self.config[provider_key] = []
        
        new_key = {
            "key": key,
            "name": name or f"{provider.title()} Key {len(self.config[provider_key]) + 1}",
            "active": True,
            "daily_limit": daily_limit,
            "usage_count": 0,
            "last_reset": datetime.now().strftime("%Y-%m-%d")
        }
        
        self.config[provider_key].append(new_key)
        self.save_config()
        
        print(f"✅ Added new {provider} API key: {new_key['name']}")
    
    def deactivate_key(self, provider: str, key_index: int):
        """Vô hiệu hóa một API key"""
        provider_key = f"{provider}_keys"
        
        if provider_key in self.config and 0 <= key_index < len(self.config[provider_key]):
            self.config[provider_key][key_index]['active'] = False
            self.save_config()
            
            key_name = self.config[provider_key][key_index].get('name', f'Key {key_index}')
            print(f"🔴 Deactivated {provider} key: {key_name}")
    
    def activate_key(self, provider: str, key_index: int):
        """Kích hoạt lại một API key"""
        provider_key = f"{provider}_keys"
        
        if provider_key in self.config and 0 <= key_index < len(self.config[provider_key]):
            self.config[provider_key][key_index]['active'] = True
            self.save_config()
            
            key_name = self.config[provider_key][key_index].get('name', f'Key {key_index}')
            print(f"🟢 Activated {provider} key: {key_name}")
    
    def get_key_status(self, provider: str) -> Dict:
        """Lấy thống kê trạng thái keys của provider"""
        provider_key = f"{provider}_keys"
        
        if provider_key not in self.config:
            return {"total": 0, "active": 0, "exhausted": 0, "available": 0}
        
        keys = self.config[provider_key]
        total = len(keys)
        active = len([k for k in keys if k.get('active', False)])
        exhausted = len([k for k in keys if self._is_key_exhausted(k)])
        available = len([k for k in keys if k.get('active', False) and not self._is_key_exhausted(k)])
        
        return {
            "total": total,
            "active": active,
            "exhausted": exhausted,
            "available": available
        }
    
    def batch_translate_with_rotation(self, texts: List[str], translator_func, provider: str = 'gemini', 
                                     max_retries: int = 3) -> List[str]:
        """
        Dịch batch texts với rotation API keys khi cần
        
        Args:
            texts (List[str]): Danh sách texts cần dịch
            translator_func: Function dịch nhận (text, api_key) -> translated_text
            provider (str): Provider name
            max_retries (int): Số lần retry khi fail
            
        Returns:
            List[str]: Danh sách texts đã dịch
        """
        results = []
        current_api_key = self.get_api_key(provider)
        
        if not current_api_key:
            print(f"❌ No available {provider} API keys for batch translation")
            return texts  # Return original texts if no keys
        
        for i, text in enumerate(texts):
            retry_count = 0
            translated = None
            
            while retry_count < max_retries and not translated:
                try:
                    # Rate limiting - thêm delay nhỏ giữa các requests
                    if i > 0:
                        time.sleep(0.5)
                    
                    translated = translator_func(text, current_api_key)
                    
                    if not translated or translated == text:
                        # API có thể bị limit, thử key khác
                        print(f"⚠️ Translation failed for text {i+1}, trying next key...")
                        current_api_key = self.get_api_key(provider)
                        
                        if not current_api_key:
                            print(f"❌ No more available {provider} keys")
                            translated = text  # Fallback to original
                            break
                        
                        retry_count += 1
                        continue
                    
                    results.append(translated)
                    break
                    
                except Exception as e:
                    print(f"❌ Error translating text {i+1}: {e}")
                    retry_count += 1
                    
                    if retry_count < max_retries:
                        # Try next API key
                        current_api_key = self.get_api_key(provider)
                        if not current_api_key:
                            break
                    
            # Nếu không dịch được sau max_retries, giữ nguyên text gốc
            if not translated:
                results.append(text)
                print(f"⚠️ Could not translate text {i+1}, keeping original")
        
        return results

# Convenience functions
def setup_api_keys():
    """Helper function để setup API keys lần đầu"""
    manager = ApiKeyManager()
    
    print("🔧 API Key Setup Helper")
    print("=" * 50)
    
    # Setup Gemini keys
    print("\n1. Gemini API Keys:")
    gemini_keys = input("Enter Gemini API keys (comma separated): ").strip()
    
    if gemini_keys:
        keys = [k.strip() for k in gemini_keys.split(',') if k.strip()]
        for i, key in enumerate(keys):
            manager.add_api_key('gemini', key, f"Gemini Account {i+1}", 1000)
    
    # Setup Google Translate keys (if any)
    print("\n2. Google Translate API Keys (optional):")
    google_keys = input("Enter Google Translate keys (comma separated, or press Enter to skip): ").strip()
    
    if google_keys:
        keys = [k.strip() for k in google_keys.split(',') if k.strip()]
        for i, key in enumerate(keys):
            manager.add_api_key('google_translate', key, f"Google Account {i+1}", 5000)
    
    print("\n✅ API Keys setup completed!")
    manager._print_key_status()
    
    return manager

if __name__ == "__main__":
    # Test the API Key Manager
    setup_api_keys()