"""
GigaChat API Tool для Telegram CLI Bot

Инструмент для работы с GigaChat API (Сбер).
Поддерживает генерацию текста, чат-сессии и различные модели.

Документация: https://developers.sber.ru/docs/ru/gigachat
"""

import os
import base64
import httpx
import asyncio
from typing import Optional, List, Dict, Any
from dataclasses import dataclass, field
from datetime import datetime, timedelta


@dataclass
class GigaChatMessage:
    """Сообщение для чата с GigaChat"""
    role: str  # 'user', 'assistant', 'system'
    content: str


@dataclass
class GigaChatConfig:
    """Конфигурация подключения к GigaChat API"""
    client_id: str
    client_secret: str
    scope: str = "GIGACHAT_API_PERS"
    auth_url: str = "https://ngw.devices.sberbank.ru:9443/api/v2/oauth"
    model: str = "GigaChat-Pro"
    api_url: str = "https://gigachat.devices.sberbank.ru/api/v2"
    timeout: int = 60


class GigaChatTool:
    """
    Инструмент для работы с GigaChat API
    
    Пример использования:
        config = GigaChatConfig(
            client_id=os.getenv("GIGACHAT_CLIENT_ID"),
            client_secret=os.getenv("GIGACHAT_CLIENT_SECRET"),
        )
        tool = GigaChatTool(config)
        
        # Простой запрос
        response = await tool.chat("Привет, как дела?")
        
        # Чат с историей
        messages = [
            GigaChatMessage(role="system", content="Ты полезный ассистент."),
            GigaChatMessage(role="user", content="Расскажи про Python"),
        ]
        response = await tool.chat(messages=messages)
    """
    
    def __init__(self, config: Optional[GigaChatConfig] = None):
        self.config = config or self._load_config_from_env()
        self._access_token: Optional[str] = None
        self._token_expires: Optional[datetime] = None
        self._chat_history: List[GigaChatMessage] = []
    
    def _load_config_from_env(self) -> GigaChatConfig:
        """Загрузка конфигурации из переменных окружения"""
        return GigaChatConfig(
            client_id=os.getenv("GIGACHAT_CLIENT_ID", ""),
            client_secret=os.getenv("GIGACHAT_CLIENT_SECRET", ""),
            scope=os.getenv("GIGACHAT_SCOPE", "GIGACHAT_API_PERS"),
            auth_url=os.getenv("GIGACHAT_AUTH_URL", "https://ngw.devices.sberbank.ru:9443/api/v2/oauth"),
            model=os.getenv("GIGACHAT_MODEL", "GigaChat-Pro"),
        )
    
    def _get_auth_headers(self) -> Dict[str, str]:
        """Получение заголовков для авторизации"""
        credentials = f"{self.config.client_id}:{self.config.client_secret}"
        encoded_credentials = base64.b64encode(credentials.encode()).decode()
        return {
            "Authorization": f"Basic {encoded_credentials}",
            "Content-Type": "application/x-www-form-urlencoded",
            "RqUID": "00000000-0000-0000-0000-000000000000",
        }
    
    async def _get_access_token(self) -> str:
        """Получение access токена для API"""
        # Проверяем кэш токена
        if self._access_token and self._token_expires:
            if datetime.now() < self._token_expires - timedelta(minutes=5):
                return self._access_token
        
        # Запрашиваем новый токен
        async with httpx.AsyncClient(verify=False) as client:
            response = await client.post(
                self.config.auth_url,
                headers=self._get_auth_headers(),
                data={"scope": self.config.scope},
                timeout=30,
            )
            response.raise_for_status()
            data = response.json()
            
            self._access_token = data["access_token"]
            # Токен действителен 30 минут, кэшируем на 25 минут
            self._token_expires = datetime.now() + timedelta(minutes=25)
            
            return self._access_token
    
    async def chat(
        self,
        messages: Optional[List[GigaChatMessage]] = None,
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
        top_p: float = 0.1,
        repetition_penalty: float = 1.0,
        use_history: bool = True,
    ) -> Dict[str, Any]:
        """
        Отправка запроса к GigaChat API
        
        Args:
            messages: Список сообщений (если None, используется история чата)
            model: Модель для генерации (если None, используется модель из конфига)
            temperature: Температура генерации (0.0 - 2.0)
            max_tokens: Максимальное количество токенов в ответе
            top_p: Параметр top-p sampling
            repetition_penalty: Штраф за повторения
            use_history: Использовать ли историю чата
        
        Returns:
            Dict с ответом API:
                - content: Текст ответа
                - model: Использованная модель
                - usage: Статистика использования токенов
                - finish_reason: Причина завершения
        """
        token = await self._get_access_token()
        
        # Формируем сообщения
        if messages is None:
            if use_history:
                messages = self._chat_history.copy()
            else:
                messages = []
        elif use_history:
            # Добавляем новые сообщения к истории
            self._chat_history.extend(messages)
            messages = self._chat_history.copy()
        
        # Преобразуем сообщения в формат API
        api_messages = [
            {"role": msg.role, "content": msg.content}
            for msg in messages
        ]
        
        payload = {
            "model": model or self.config.model,
            "messages": api_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "top_p": top_p,
            "repetition_penalty": repetition_penalty,
        }
        
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
        
        async with httpx.AsyncClient(verify=False) as client:
            response = await client.post(
                f"{self.config.api_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=self.config.timeout,
            )
            response.raise_for_status()
            data = response.json()
        
        # Добавляем ответ ассистента в историю
        if use_history and data.get("choices"):
            assistant_message = data["choices"][0]["message"]
            self._chat_history.append(GigaChatMessage(
                role=assistant_message["role"],
                content=assistant_message["content"],
            ))
        
        return {
            "content": data["choices"][0]["message"]["content"] if data.get("choices") else "",
            "model": data.get("model", self.config.model),
            "usage": data.get("usage", {}),
            "finish_reason": data["choices"][0]["finish_reason"] if data.get("choices") else "",
        }
    
    def clear_history(self):
        """Очистка истории чата"""
        self._chat_history = []
    
    def get_history(self) -> List[GigaChatMessage]:
        """Получение истории чата"""
        return self._chat_history.copy()
    
    def set_system_prompt(self, prompt: str):
        """Установка системного промпта (добавляется в начало истории)"""
        # Удаляем старый системный промпт если есть
        self._chat_history = [
            msg for msg in self._chat_history if msg.role != "system"
        ]
        # Добавляем новый в начало
        self._chat_history.insert(0, GigaChatMessage(role="system", content=prompt))
    
    async def generate_image(
        self,
        prompt: str,
        model: str = " Kandinsky-2",
        size: str = "1024x1024",
    ) -> Dict[str, Any]:
        """
        Генерация изображений через GigaChat (Kandinsky)
        
        Args:
            prompt: Текстовое описание изображения
            model: Модель для генерации
            size: Размер изображения
        
        Returns:
            Dict с результатом генерации
        """
        token = await self._get_access_token()
        
        payload = {
            "model": model,
            "prompt": prompt,
            "size": size,
        }
        
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
        
        async with httpx.AsyncClient(verify=False) as client:
            # Запуск генерации
            response = await client.post(
                f"{self.config.api_url}/images/generations",
                headers=headers,
                json=payload,
                timeout=self.config.timeout,
            )
            response.raise_for_status()
            data = response.json()
        
        return data
    
    async def get_models(self) -> List[str]:
        """Получение списка доступных моделей"""
        token = await self._get_access_token()
        
        headers = {
            "Authorization": f"Bearer {token}",
        }
        
        async with httpx.AsyncClient(verify=False) as client:
            response = await client.get(
                f"{self.config.api_url}/models",
                headers=headers,
                timeout=30,
            )
            response.raise_for_status()
            data = response.json()
        
        return [model["id"] for model in data.get("data", [])]


# Утилита для создания инструмента в формате бота
def create_gigachat_tool():
    """
    Создает экземпляр GigaChatTool с конфигурацией из окружения
    
    Returns:
        GigaChatTool или None если конфигурация не задана
    """
    if not os.getenv("GIGACHAT_CLIENT_ID") or not os.getenv("GIGACHAT_CLIENT_SECRET"):
        return None
    
    return GigaChatTool()


if __name__ == "__main__":
    # Пример использования
    async def main():
        tool = create_gigachat_tool()
        if not tool:
            print("GigaChat не настроен. Проверьте переменные окружения.")
            return
        
        # Простой запрос
        response = await tool.chat("Привет! Расскажи кратко про себя.")
        print(f"Ответ: {response['content']}")
        print(f"Модель: {response['model']}")
        print(f"Токены: {response['usage']}")
    
    asyncio.run(main())

# ===========================================
# Интеграция с реестром инструментов бота
# ===========================================

from bot.tools import BaseTool, ToolResult, register_tool


@register_tool
class GigaChatCapability(BaseTool):
    """
    Capability-обёртка для GigaChat API.
    
    Позволяет использовать GigaChat через реестр инструментов бота.
    """
    
    name = "gigachat"
    description = "Генерация ответов AI через GigaChat API (Сбер). Альтернатива Qwen Code."
    category = "ai"
    
    def __init__(self):
        self._provider = None
    
    def _get_provider(self):
        """Ленивая инициализация провайдера"""
        if self._provider is None:
            from qwen_integration import GigaChatProvider
            self._provider = GigaChatProvider()
        return self._provider
    
    async def execute(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
        **kwargs
    ) -> ToolResult:
        """
        Выполнить запрос к GigaChat API.
        
        Args:
            prompt: Запрос пользователя
            system_prompt: Системный промпт (роль ассистента)
            temperature: Температура генерации (0.0-2.0)
            max_tokens: Максимум токенов в ответе
        """
        provider = self._get_provider()
        
        if not provider.is_available():
            return ToolResult(
                success=False,
                error=provider.get_error() or "GigaChat не доступен",
            )
        
        result = await provider.chat(
            prompt=prompt,
            system_prompt=system_prompt,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        
        if result.get("success"):
            return ToolResult(
                success=True,
                data={
                    "content": result.get("content", ""),
                    "model": result.get("model", "GigaChat-Pro"),
                    "usage": result.get("usage", {}),
                },
                metadata={
                    "model": result.get("model"),
                    "tokens": result.get("usage"),
                }
            )
        else:
            return ToolResult(
                success=False,
                error=result.get("error", "Неизвестная ошибка GigaChat"),
            )