#!/usr/bin/env python3
"""
Скрипт для добавления новых LLM моделей
"""

import requests
import json
from typing import Dict, List

# Конфигурация
API_BASE = "http://localhost:8888"

def add_model_to_config(model_name: str, model_id: str, description: str = ""):
    """
    Добавить модель в конфигурацию llm_config.py
    
    Args:
        model_name: Человекочитаемое имя (например, 'fast', 'smart')
        model_id: ID модели в API (например, 'gpt-4o-mini')
        description: Описание модели
    """
    
    config_file = "/root/engine/public_oversight/hotels/llm_config.py"
    
    # Читаем текущий конфиг
    with open(config_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Находим секцию models и добавляем новую модель
    if "'models': {" in content:
        # Ищем место для вставки
        models_start = content.find("'models': {")
        models_end = content.find("},", models_start) + 2
        
        # Формируем новую строку
        new_model_line = f"        '{model_name}': '{model_id}',          # {description}\n"
        
        # Вставляем перед закрывающей скобкой
        new_content = content[:models_end-2] + new_model_line + content[models_end-2:]
        
        # Записываем обратно
        with open(config_file, 'w', encoding='utf-8') as f:
            f.write(new_content)
        
        print(f"✅ Модель '{model_name}' добавлена в конфиг")
        return True
    else:
        print("❌ Не найдена секция models в конфиге")
        return False


def test_model(model_id: str) -> Dict:
    """
    Протестировать модель
    
    Args:
        model_id: ID модели для тестирования
    
    Returns:
        Результат теста
    """
    
    # Переключаемся на модель
    switch_response = requests.post(
        f"{API_BASE}/api/llm/switch",
        json={"model": model_id}
    )
    
    if switch_response.status_code != 200:
        return {"error": f"Ошибка переключения: {switch_response.status_code}"}
    
    # Тестируем модель
    test_prompt = "Сколько отелей в Чукотском автономном округе?"
    
    start_time = requests.utils.time.time()
    chat_response = requests.post(
        f"{API_BASE}/api/chat",
        json={"message": test_prompt}
    )
    end_time = requests.utils.time.time()
    
    if chat_response.status_code == 200:
        data = chat_response.json()
        return {
            "model": model_id,
            "response_time": round((end_time - start_time) * 1000, 2),
            "response_length": len(data['response']),
            "response_preview": data['response'][:200] + "...",
            "success": True
        }
    else:
        return {"error": f"Ошибка чата: {chat_response.status_code}"}


def list_available_models() -> List[str]:
    """Получить список доступных моделей"""
    
    response = requests.get(f"{API_BASE}/api/llm/models")
    if response.status_code == 200:
        data = response.json()
        return list(data['models'].values())
    else:
        return []


def add_custom_model():
    """Интерактивное добавление модели"""
    
    print("=" * 60)
    print("🤖 ДОБАВЛЕНИЕ НОВОЙ LLM МОДЕЛИ")
    print("=" * 60)
    
    # Получаем информацию о модели
    model_name = input("Введите имя модели (например, 'custom'): ").strip()
    model_id = input("Введите ID модели (например, 'gpt-4o-mini'): ").strip()
    description = input("Введите описание (опционально): ").strip()
    
    if not model_name or not model_id:
        print("❌ Имя и ID модели обязательны!")
        return
    
    print(f"\n📋 Добавляем модель:")
    print(f"  Имя: {model_name}")
    print(f"  ID: {model_id}")
    print(f"  Описание: {description}")
    
    confirm = input("\nПродолжить? (y/N): ").strip().lower()
    if confirm != 'y':
        print("❌ Отменено")
        return
    
    # Добавляем в конфиг
    if add_model_to_config(model_name, model_id, description):
        print(f"\n🧪 Тестируем модель {model_id}...")
        
        # Тестируем
        result = test_model(model_id)
        
        if result.get('success'):
            print(f"✅ Модель работает!")
            print(f"  Время ответа: {result['response_time']}ms")
            print(f"  Длина ответа: {result['response_length']} символов")
            print(f"  Превью: {result['response_preview']}")
        else:
            print(f"❌ Ошибка тестирования: {result.get('error')}")
        
        print(f"\n🔄 Перезапустите веб-интерфейс для применения изменений:")
        print(f"  pkill -f web_interface.py")
        print(f"  cd /root/engine/public_oversight/hotels")
        print(f"  source venv/bin/activate")
        print(f"  python web_interface.py &")


def quick_add_popular_models():
    """Быстрое добавление популярных моделей"""
    
    popular_models = [
        ("gpt35", "gpt-3.5-turbo", "GPT-3.5 Turbo (классика)"),
        ("gpt4v", "gpt-4-vision-preview", "GPT-4 Vision (анализ изображений)"),
        ("claude", "claude-3-haiku", "Claude 3 Haiku (через OpenRouter)"),
        ("gemini", "gemini-pro", "Gemini Pro (через OpenRouter)"),
    ]
    
    print("🚀 Быстрое добавление популярных моделей:")
    
    for i, (name, model_id, desc) in enumerate(popular_models, 1):
        print(f"{i}. {desc}")
    
    choice = input("\nВыберите номер модели (1-4) или 0 для выхода: ").strip()
    
    try:
        choice_idx = int(choice) - 1
        if 0 <= choice_idx < len(popular_models):
            name, model_id, desc = popular_models[choice_idx]
            add_model_to_config(name, model_id, desc)
            print(f"✅ Модель {desc} добавлена!")
        else:
            print("❌ Неверный выбор")
    except ValueError:
        print("❌ Введите число")


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        command = sys.argv[1]
        
        if command == "list":
            print("📋 Доступные модели:")
            models = list_available_models()
            for model in models:
                print(f"  - {model}")
        
        elif command == "test":
            if len(sys.argv) > 2:
                model_id = sys.argv[2]
                print(f"🧪 Тестируем модель: {model_id}")
                result = test_model(model_id)
                print(json.dumps(result, indent=2, ensure_ascii=False))
            else:
                print("❌ Укажите ID модели для тестирования")
        
        elif command == "add":
            if len(sys.argv) > 3:
                name, model_id = sys.argv[2], sys.argv[3]
                description = sys.argv[4] if len(sys.argv) > 4 else ""
                add_model_to_config(name, model_id, description)
            else:
                print("❌ Использование: python add_model.py add <name> <model_id> [description]")
        
        else:
            print("❌ Неизвестная команда")
    
    else:
        # Интерактивный режим
        print("🤖 УПРАВЛЕНИЕ LLM МОДЕЛЯМИ")
        print("1. Добавить модель вручную")
        print("2. Быстрое добавление популярных моделей")
        print("3. Список доступных моделей")
        print("4. Тест модели")
        
        choice = input("\nВыберите действие (1-4): ").strip()
        
        if choice == "1":
            add_custom_model()
        elif choice == "2":
            quick_add_popular_models()
        elif choice == "3":
            models = list_available_models()
            print("📋 Доступные модели:")
            for model in models:
                print(f"  - {model}")
        elif choice == "4":
            model_id = input("Введите ID модели для тестирования: ").strip()
            if model_id:
                result = test_model(model_id)
                print(json.dumps(result, indent=2, ensure_ascii=False))
        else:
            print("❌ Неверный выбор")