#!/usr/bin/env python3
"""
API для получения списка моделей от разных провайдеров
"""

import requests
import json
from typing import Dict, List, Optional
from llm_config import OPENAI_CONFIG, OPENROUTER_CONFIG, OLLAMA_CONFIG

class ModelProvider:
    """Базовый класс для провайдеров моделей"""
    
    def __init__(self, config: Dict):
        self.config = config
        self.provider = config['provider']
        self.api_base = config['api_base']
        self.api_key = config.get('api_key')
        self.proxy = config.get('proxy')
    
    def get_models(self) -> List[Dict]:
        """Получить список доступных моделей"""
        raise NotImplementedError
    
    def _make_request(self, endpoint: str, method: str = 'GET', data: Optional[Dict] = None) -> Optional[Dict]:
        """Универсальный запрос к API"""
        url = f"{self.api_base}/{endpoint}"
        
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        
        if self.provider == 'openrouter':
            headers["HTTP-Referer"] = "https://hotel-audit.ru"
            headers["X-Title"] = "Hotel Audit System"
        
        headers["Content-Type"] = "application/json"
        
        proxies = None
        if self.proxy:
            proxies = {'http': self.proxy, 'https': self.proxy}
        
        try:
            if method == 'GET':
                response = requests.get(url, headers=headers, proxies=proxies, timeout=30)
            elif method == 'POST':
                response = requests.post(url, headers=headers, json=data, proxies=proxies, timeout=30)
            
            if response.status_code == 200:
                return response.json()
            else:
                print(f"Ошибка {self.provider}: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            print(f"Ошибка соединения {self.provider}: {e}")
            return None


class OpenAIProvider(ModelProvider):
    """Провайдер OpenAI"""
    
    def get_models(self) -> List[Dict]:
        """Получить модели OpenAI"""
        data = self._make_request("models")
        if not data:
            return []
        
        models = []
        for model in data.get('data', []):
            model_id = model['id']
            
            # Фильтруем только чат-модели
            if any(x in model_id.lower() for x in ['gpt-3.5', 'gpt-4', 'o1']):
                models.append({
                    'id': model_id,
                    'name': model_id,
                    'provider': 'openai',
                    'description': self._get_model_description(model_id),
                    'context_length': model.get('context_length', 0),
                    'pricing': self._get_pricing(model_id)
                })
        
        return sorted(models, key=lambda x: x['name'])
    
    def _get_model_description(self, model_id: str) -> str:
        """Описание модели"""
        descriptions = {
            'gpt-4o-mini': 'Быстрая и дешёвая модель',
            'gpt-4o': 'Самая умная модель OpenAI',
            'gpt-4-turbo': 'Баланс скорости и качества',
            'gpt-4': 'Классическая GPT-4',
            'gpt-3.5-turbo': 'Проверенная временем модель',
            'o1-preview': 'Новая рассуждающая модель',
            'o1-mini': 'Компактная рассуждающая модель'
        }
        
        for key, desc in descriptions.items():
            if key in model_id.lower():
                return desc
        
        return 'Модель OpenAI'
    
    def _get_pricing(self, model_id: str) -> Dict:
        """Примерная стоимость"""
        pricing = {
            'gpt-4o-mini': {'input': 0.15, 'output': 0.60},
            'gpt-4o': {'input': 5.00, 'output': 15.00},
            'gpt-4-turbo': {'input': 10.00, 'output': 30.00},
            'gpt-3.5-turbo': {'input': 0.50, 'output': 1.50},
            'o1-preview': {'input': 15.00, 'output': 60.00},
            'o1-mini': {'input': 3.00, 'output': 12.00},
            'gpt-4': {'input': 30.00, 'output': 60.00},
            'gpt-3.5': {'input': 0.50, 'output': 1.50}
        }
        
        # Ищем точное совпадение или частичное
        for key, price in pricing.items():
            if key in model_id.lower():
                return price
        
        # Если не найдено, возвращаем базовую цену
        return {'input': 0.50, 'output': 1.50}


class OpenRouterProvider(ModelProvider):
    """Провайдер OpenRouter"""
    
    def get_models(self) -> List[Dict]:
        """Получить модели OpenRouter"""
        data = self._make_request("models")
        if not data:
            return []
        
        models = []
        for model in data.get('data', []):
            model_id = model['id']
            
            # Фильтруем популярные модели
            if any(x in model_id.lower() for x in ['claude', 'gemini', 'llama', 'mistral', 'qwen']):
                models.append({
                    'id': model_id,
                    'name': model_id,
                    'provider': 'openrouter',
                    'description': self._get_model_description(model_id),
                    'context_length': model.get('context_length', 0),
                    'pricing': self._get_pricing(model_id)
                })
        
        return sorted(models, key=lambda x: x['name'])
    
    def _get_model_description(self, model_id: str) -> str:
        """Описание модели"""
        if 'claude' in model_id.lower():
            if 'haiku' in model_id.lower():
                return 'Claude 3 Haiku - быстрая модель Anthropic'
            elif 'sonnet' in model_id.lower():
                return 'Claude 3.5 Sonnet - умная модель Anthropic'
            elif 'opus' in model_id.lower():
                return 'Claude 3 Opus - самая мощная модель Anthropic'
            else:
                return 'Claude модель от Anthropic'
        
        elif 'gemini' in model_id.lower():
            return 'Gemini модель от Google'
        
        elif 'llama' in model_id.lower():
            return 'Llama модель от Meta'
        
        elif 'mistral' in model_id.lower():
            return 'Mistral модель'
        
        elif 'qwen' in model_id.lower():
            return 'Qwen модель от Alibaba'
        
        return 'Модель через OpenRouter'
    
    def _get_pricing(self, model_id: str) -> Dict:
        """Примерная стоимость"""
        pricing = {
            'claude-3-haiku': {'input': 0.25, 'output': 1.25},
            'claude-3.5-sonnet': {'input': 3.00, 'output': 15.00},
            'claude-3-opus': {'input': 15.00, 'output': 75.00},
            'gemini-flash': {'input': 0.075, 'output': 0.30},
            'gemini-pro': {'input': 0.50, 'output': 1.50},
            'llama3.1': {'input': 0.20, 'output': 0.20},
            'mistral': {'input': 0.25, 'output': 0.25}
        }
        
        for key, price in pricing.items():
            if key in model_id.lower():
                return price
        
        return {'input': 0, 'output': 0}


class OllamaProvider(ModelProvider):
    """Провайдер Ollama (локальный)"""
    
    def get_models(self) -> List[Dict]:
        """Получить модели Ollama"""
        data = self._make_request("api/tags")
        if not data:
            return []
        
        models = []
        for model in data.get('models', []):
            model_id = model['name']
            
            models.append({
                'id': model_id,
                'name': model_id,
                'provider': 'ollama',
                'description': self._get_model_description(model_id),
                'context_length': model.get('size', 0),
                'pricing': {'input': 0, 'output': 0}  # Бесплатно
            })
        
        return sorted(models, key=lambda x: x['name'])
    
    def _get_model_description(self, model_id: str) -> str:
        """Описание модели"""
        if 'llama' in model_id.lower():
            return 'Llama модель (локальная)'
        elif 'mistral' in model_id.lower():
            return 'Mistral модель (локальная)'
        elif 'qwen' in model_id.lower():
            return 'Qwen модель (локальная)'
        elif 'codellama' in model_id.lower():
            return 'Code Llama для программирования'
        else:
            return 'Локальная модель Ollama'


def get_all_models() -> Dict[str, List[Dict]]:
    """Получить все модели от всех провайдеров"""
    
    providers = {
        'openai': OpenAIProvider(OPENAI_CONFIG),
        'openrouter': OpenRouterProvider(OPENROUTER_CONFIG),
        'ollama': OllamaProvider(OLLAMA_CONFIG)
    }
    
    all_models = {}
    
    for provider_name, provider in providers.items():
        print(f"Загружаю модели {provider_name}...")
        try:
            models = provider.get_models()
            all_models[provider_name] = models
            print(f"✅ {provider_name}: {len(models)} моделей")
        except Exception as e:
            print(f"❌ {provider_name}: ошибка - {e}")
            all_models[provider_name] = []
    
    return all_models


if __name__ == "__main__":
    print("=" * 70)
    print("🤖 ЗАГРУЗКА МОДЕЛЕЙ ОТ ПРОВАЙДЕРОВ")
    print("=" * 70)
    
    all_models = get_all_models()
    
    print(f"\n📊 ИТОГО:")
    for provider, models in all_models.items():
        print(f"  {provider.upper()}: {len(models)} моделей")
    
    print(f"\n🎯 ПРИМЕРЫ МОДЕЛЕЙ:")
    for provider, models in all_models.items():
        if models:
            print(f"\n{provider.upper()}:")
            for model in models[:3]:  # Показываем первые 3
                print(f"  - {model['name']}: {model['description']}")
    
    print("=" * 70)