# -*- coding: utf-8 -*-
"""
HTTP-клиент для взаимодействия с RAG-сервером.
Используется ботами вместо локального RAGAPI, когда RAG-ядро вынесено
в отдельный HTTP-сервис.

Клиент отправляет асинхронные HTTP-запросы на эндпоинты /rag/query и /rag/index.
Все методы повторяют сигнатуру RAGAPI для бесшовной замены.

ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.

Улучшена обработка сетевых ошибок с повторными попытками.

ДОБАВЛЕНО:
- Метод vision() для распознавания текста на изображениях
- Метод transcribe() для транскрибации аудио
- Аутентификация через API-ключ (заголовок X-API-Key)
"""

import os
import logging
from typing import Optional, Dict, List, Any
import aiohttp
import asyncio

logger = logging.getLogger(__name__)


class RAGClient:
    """
    HTTP-клиент для RAG-сервера с поддержкой повторных попыток при сетевых ошибках.
    """

    def __init__(
        self,
        server_url: str,
        api_key: Optional[str] = None,  # <-- добавлен API-ключ
        timeout: int = 60,
        max_retries: int = 3,
        retry_delay: float = 1.0
    ):
        """
        Инициализация клиента.

        Аргументы:
            server_url (str): базовый URL RAG-сервера (например, http://localhost:8080)
            api_key (Optional[str]): API-ключ для аутентификации (передаётся в заголовке X-API-Key)
            timeout (int): таймаут на запрос в секундах
            max_retries (int): максимальное количество повторных попыток при ошибке
            retry_delay (float): базовая задержка между попытками (секунды)
        """
        self.server_url = server_url.rstrip('/')
        self.api_key = api_key              # <-- сохраняем ключ
        self.timeout = timeout
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self._session: Optional[aiohttp.ClientSession] = None
        logger.info(
            f"RAGClient инициализирован с сервером {server_url}, "
            f"таймаут={timeout}с, макс. попыток={max_retries}, "
            f"API-ключ {'задан' if api_key else 'не задан'}"
        )

    async def _get_session(self) -> aiohttp.ClientSession:
        """
        Возвращает сессию aiohttp (создаёт при первом вызове).
        Сессия переиспользуется для всех запросов для эффективности.
        """
        if self._session is None:
            # Создаём сессию с таймаутом, который применяется ко всем запросам
            timeout_obj = aiohttp.ClientTimeout(total=self.timeout)
            self._session = aiohttp.ClientSession(timeout=timeout_obj)
            logger.debug("HTTP-сессия создана")
        return self._session

    async def close(self):
        """Закрывает HTTP-сессию."""
        if self._session:
            await self._session.close()
            self._session = None
            logger.debug("RAGClient сессия закрыта")

    async def _post(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Внутренний метод для отправки POST-запроса с повторными попытками.
        Автоматически добавляет заголовок X-API-Key, если ключ задан.

        Аргументы:
            endpoint (str): путь эндпоинта (например, '/rag/query')
            payload (Dict): тело запроса в формате JSON

        Возвращает:
            Dict: ответ сервера

        Исключения:
            aiohttp.ClientError: если все попытки не удались
            asyncio.TimeoutError: если все попытки истекли по таймауту
        """
        url = f"{self.server_url}{endpoint}"
        session = await self._get_session()

        # Подготавливаем заголовки
        headers = {}
        if self.api_key:
            headers['X-API-Key'] = self.api_key

        last_exception = None

        for attempt in range(self.max_retries):
            try:
                logger.debug(f"Попытка {attempt + 1}/{self.max_retries} к {endpoint}")
                async with session.post(url, json=payload, headers=headers) as resp:
                    if resp.status == 200:
                        return await resp.json()
                    else:
                        # Обрабатываем HTTP ошибки (4xx, 5xx)
                        error_text = await resp.text()
                        error_msg = f"HTTP {resp.status}: {error_text[:200]}"
                        logger.error(error_msg)
                        # Для 5xx ошибок делаем повтор, для 4xx – нет (это ошибка клиента)
                        if 500 <= resp.status < 600:
                            raise aiohttp.ClientResponseError(
                                status=resp.status,
                                message=error_msg,
                                headers=resp.headers
                            )
                        else:
                            # 4xx – не повторяем, сразу выбрасываем
                            raise aiohttp.ClientResponseError(
                                status=resp.status,
                                message=error_msg,
                                headers=resp.headers
                            )

            except (aiohttp.ClientConnectorError, aiohttp.ServerDisconnectedError) as e:
                # Ошибки подключения или разрыва соединения – повторяем с увеличенной задержкой
                logger.warning(
                    f"Сетевая ошибка при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
                )
                last_exception = e
                if attempt < self.max_retries - 1:
                    # Для сетевых ошибок задержка больше: 2, 4, 8 секунд
                    wait = (self.retry_delay * 2) * (2 ** attempt)
                    logger.info(f"Повтор через {wait:.1f} секунд")
                    await asyncio.sleep(wait)
                else:
                    logger.error(f"Все попытки к {endpoint} не удались из-за сетевых проблем")

            except (aiohttp.ClientPayloadError, aiohttp.ClientResponseError) as e:
                # Ошибки ответа (например, повреждённые данные)
                logger.warning(
                    f"Ошибка ответа при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
                )
                last_exception = e
                if attempt < self.max_retries - 1:
                    wait = self.retry_delay * (2 ** attempt)
                    logger.info(f"Повтор через {wait:.1f} секунд")
                    await asyncio.sleep(wait)
                else:
                    logger.error(f"Все попытки к {endpoint} не удались из-за ошибок ответа")

            except asyncio.TimeoutError as e:
                # Таймаут
                logger.warning(
                    f"Таймаут при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
                )
                last_exception = e
                if attempt < self.max_retries - 1:
                    wait = self.retry_delay * (2 ** attempt)
                    logger.info(f"Повтор через {wait:.1f} секунд")
                    await asyncio.sleep(wait)
                else:
                    logger.error(f"Все попытки к {endpoint} не удались из-за таймаута")

            except aiohttp.ClientError as e:
                # Все остальные ошибки клиента (общий класс)
                logger.warning(
                    f"Ошибка клиента при попытке {attempt + 1}/{self.max_retries} к {endpoint}: {e}"
                )
                last_exception = e
                if attempt < self.max_retries - 1:
                    wait = self.retry_delay * (2 ** attempt)
                    logger.info(f"Повтор через {wait:.1f} секунд")
                    await asyncio.sleep(wait)
                else:
                    logger.error(f"Все попытки к {endpoint} не удались")

            except Exception as e:
                # Неожиданная ошибка – не повторяем, логируем и выбрасываем
                logger.error(f"Неожиданная ошибка при запросе к {endpoint}: {e}", exc_info=True)
                raise

        # Если мы вышли из цикла, значит, все попытки провалились
        error_msg = f"Не удалось выполнить запрос к {endpoint} после {self.max_retries} попыток"
        if last_exception:
            error_msg += f": {last_exception}"
        logger.error(error_msg)
        raise aiohttp.ClientError(error_msg)

    async def health_check(self) -> bool:
        """
        Проверяет доступность RAG-сервера через эндпоинт /health.
        Возвращает True, если сервер отвечает со статусом 200.
        """
        try:
            session = await self._get_session()
            headers = {}
            if self.api_key:
                headers['X-API-Key'] = self.api_key
            async with session.get(f"{self.server_url}/health", timeout=10, headers=headers) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    if data.get('status') == 'healthy':
                        logger.debug("RAG-сервер доступен (health check пройден)")
                        return True
                    else:
                        logger.warning(f"RAG-сервер вернул нездоровый статус: {data}")
                        return False
                else:
                    logger.warning(f"Health check вернул HTTP {resp.status}")
                    return False
        except Exception as e:
            logger.error(f"Ошибка при health check: {e}")
            return False

    async def query(
        self,
        query: str,
        user_jid: str,
        room_jid: Optional[str],
        prompts: Optional[Dict[str, str]] = None,
        intent_override: Optional[str] = None,
        last_file_path: Optional[str] = None,
        last_file_text: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Выполняет RAG-запрос через HTTP-вызов /rag/query.

        Сигнатура полностью совпадает с RAGAPI.query для бесшовной замены.
        ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
        """
        payload = {
            "query": query,
            "user_jid": user_jid,
            "room_jid": room_jid,
            "prompts": prompts or {},
            "intent_override": intent_override,
            "last_file_path": last_file_path,
            "last_file_text": last_file_text,
        }
        # Убираем None-значения, чтобы не засорять запрос
        payload = {k: v for k, v in payload.items() if v is not None}

        try:
            result = await self._post("/rag/query", payload)
            logger.debug(f"RAGClient.query успешно выполнен для {user_jid}")
            return result
        except Exception as e:
            logger.error(f"Ошибка в RAGClient.query: {e}", exc_info=True)
            # Возвращаем структурированный ответ с ошибкой
            return {
                "answer": f"⚠️ Ошибка связи с RAG-сервером: {str(e)}",
                "intent": "ERROR",
                "context": "",
                "sources": [],
                "confidence": None,
                "error": str(e)
            }

    async def index_document(
        self,
        file_name: str,
        file_text: str,
        user_jid: str,
        room_jid: Optional[str],
        is_global: bool = False,
        title: Optional[str] = None,
        metadata: Optional[Dict] = None,
        file_hash: Optional[str] = None,
        update_if_exists: bool = True
    ) -> Dict[str, Any]:
        """
        Индексирует документ через HTTP-вызов /rag/index.

        Сигнатура полностью совпадает с RAGAPI.index_document.
        """
        payload = {
            "file_name": file_name,
            "file_text": file_text,
            "user_jid": user_jid,
            "room_jid": room_jid,
            "is_global": is_global,
            "title": title,
            "metadata": metadata,
            "file_hash": file_hash,
            "update_if_exists": update_if_exists,
        }
        payload = {k: v for k, v in payload.items() if v is not None}

        try:
            result = await self._post("/rag/index", payload)
            logger.info(f"Документ {file_name} проиндексирован через RAG-сервер")
            return result
        except Exception as e:
            logger.error(f"Ошибка в RAGClient.index_document: {e}", exc_info=True)
            return {"doc_id": None, "chunk_count": 0, "error": str(e)}

    # ==========================================================
    # НОВЫЕ МЕТОДЫ: OCR и транскрибация
    # ==========================================================

    async def vision(self, file_path: str) -> Dict[str, Any]:
        """
        Отправляет изображение на RAG-сервер для распознавания текста (OCR).
        
        Аргументы:
            file_path: путь к файлу изображения
            
        Возвращает:
            Dict с ключами: text (str), status (str), error (опционально)
        """
        if not os.path.exists(file_path):
            return {"text": "", "status": "file_not_found", "error": "Файл не найден"}
        
        session = await self._get_session()
        url = f"{self.server_url}/rag/vision"
        
        # Добавляем заголовок с ключом
        headers = {}
        if self.api_key:
            headers['X-API-Key'] = self.api_key
        
        try:
            with open(file_path, 'rb') as f:
                files = {'file': (os.path.basename(file_path), f, 'image/jpeg')}
                async with session.post(url, data=files, headers=headers) as resp:
                    if resp.status == 200:
                        return await resp.json()
                    else:
                        error_text = await resp.text()
                        logger.error(f"OCR ошибка: {resp.status} - {error_text[:200]}")
                        return {"text": "", "status": "error", "error": error_text}
        except Exception as e:
            logger.error(f"Ошибка вызова OCR: {e}", exc_info=True)
            return {"text": "", "status": "error", "error": str(e)}

    async def transcribe(self, file_path: str) -> Dict[str, Any]:
        """
        Отправляет аудиофайл на RAG-сервер для транскрибации (SaluteSpeech).
        
        Аргументы:
            file_path: путь к аудиофайлу
            
        Возвращает:
            Dict с ключами: text (str), status (str), error (опционально)
        """
        if not os.path.exists(file_path):
            return {"text": "", "status": "file_not_found", "error": "Файл не найден"}
        
        session = await self._get_session()
        url = f"{self.server_url}/rag/transcribe"
        
        headers = {}
        if self.api_key:
            headers['X-API-Key'] = self.api_key
        
        try:
            with open(file_path, 'rb') as f:
                files = {'file': (os.path.basename(file_path), f, 'audio/ogg')}
                async with session.post(url, data=files, headers=headers) as resp:
                    if resp.status == 200:
                        return await resp.json()
                    else:
                        error_text = await resp.text()
                        logger.error(f"Транскрибация ошибка: {resp.status} - {error_text[:200]}")
                        return {"text": "", "status": "error", "error": error_text}
        except Exception as e:
            logger.error(f"Ошибка вызова транскрибации: {e}", exc_info=True)
            return {"text": "", "status": "error", "error": str(e)}