#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RAG-сервер – отдельный HTTP-сервис для обработки RAG-запросов.
Предоставляет API для выполнения RAG-запросов и индексации документов.
ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL).

ДОБАВЛЕНО:
- Эндпоинт /rag/vision для распознавания текста на изображениях (OCR)
- Эндпоинт /rag/transcribe для транскрибации аудио (SaluteSpeech)
- Использование Pydantic-модели AppConfig для конфигурации.
- Аутентификация через API-ключ (X-API-Key)
- Полная OpenAPI-документация с примерами и тегами
"""

import logging
import os
import sys
import tempfile
import shutil
import asyncio
from pathlib import Path
from typing import Optional, Dict, List, Any

# Добавляем путь к rag, если запускаем из корня проекта
sys.path.insert(0, str(Path(__file__).parent.parent))

from fastapi import FastAPI, HTTPException, Depends, UploadFile, File as FastAPIFile, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.openapi.utils import get_openapi
from pydantic import BaseModel, Field, HttpUrl, constr
import uvicorn

# Импорты наших модулей
from .utils.config_loader import load_config
from .utils.logger import setup_logging
from .services.postgres_service import PostgresService
from .services.qdrant_service import QdrantService
from .services.embedding_service import EmbeddingService
from .services.kb_service import KBService
from .services.giga_client import GigaClient
from .services.file_service import FileService
from .rag_orchestrator import RAGOrchestrator
from .auth import verify_api_key, set_auth_config

logger = logging.getLogger(__name__)


# ============================================================
# МОДЕЛИ ДАННЫХ (Pydantic) ДЛЯ API
# ============================================================

class QueryRequest(BaseModel):
    """
    Модель запроса к RAG-серверу на выполнение RAG-запроса.
    ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
    """
    query: str = Field(
        ..., 
        description="Текст запроса пользователя",
        example="Какое тактовое время установлено на линии сборки?"
    )
    user_jid: str = Field(
        ..., 
        description="JID пользователя (без ресурса)",
        example="user@domain.ru"
    )
    room_jid: Optional[str] = Field(
        None, 
        description="JID комнаты (None для личного чата)",
        example="room@conference.domain.ru"
    )
    prompts: Dict[str, str] = Field(
        default_factory=dict,
        description="Словарь промптов (system, synthesis, intent, expand, critique и т.д.)"
    )
    intent_override: Optional[str] = Field(
        None, 
        description="Принудительное переопределение намерения",
        example="FACT"
    )
    last_file_path: Optional[str] = Field(
        None, 
        description="Путь к последнему загруженному файлу на сервере"
    )
    last_file_text: Optional[str] = Field(
        None, 
        description="Текст последнего загруженного файла"
    )

    class Config:
        schema_extra = {
            "example": {
                "query": "Какое тактовое время установлено на линии сборки?",
                "user_jid": "user@domain.ru",
                "room_jid": None,
                "prompts": {},
                "intent_override": None,
                "last_file_path": None,
                "last_file_text": None
            }
        }


class QueryResponse(BaseModel):
    """Модель ответа RAG-сервера на RAG-запрос."""
    answer: str = Field(
        ..., 
        description="Итоговый ответ",
        example="Тактовое время на линии сборки составляет 120 секунд. [Источник: Регламент_линии_2025.docx]"
    )
    intent: str = Field(
        default="GENERAL", 
        description="Распознанное намерение",
        example="FACT"
    )
    context: Optional[str] = Field(
        None, 
        description="Использованный контекст (для отладки)"
    )
    sources: List[str] = Field(
        default_factory=list, 
        description="Список источников",
        example=["Регламент_линии_2025.docx"]
    )
    confidence: Optional[float] = Field(
        None, 
        description="Оценка уверенности (0-1)",
        example=0.95
    )
    error: Optional[str] = Field(
        None, 
        description="Сообщение об ошибке, если есть"
    )


class IndexRequest(BaseModel):
    """Модель запроса на индексацию документа."""
    file_name: str = Field(
        ..., 
        description="Исходное имя файла",
        example="Регламент_отдела.docx"
    )
    file_text: str = Field(
        ..., 
        description="Извлечённый текст документа",
        example="Полный текст документа..."
    )
    user_jid: str = Field(
        ..., 
        description="JID владельца документа",
        example="user@domain.ru"
    )
    room_jid: Optional[str] = Field(
        None, 
        description="JID комнаты (если документ комнатный)",
        example="room@conference.domain.ru"
    )
    is_global: bool = Field(
        False, 
        description="Глобальный ли документ"
    )
    title: Optional[str] = Field(
        None, 
        description="Отображаемое название"
    )
    metadata: Optional[Dict[str, Any]] = Field(
        None, 
        description="Дополнительные метаданные"
    )
    file_hash: Optional[str] = Field(
        None, 
        description="SHA-256 хеш содержимого"
    )
    update_if_exists: bool = Field(
        True, 
        description="Заменять ли существующий документ"
    )


class IndexResponse(BaseModel):
    """Модель ответа на индексацию."""
    doc_id: Optional[int] = Field(
        None, 
        description="Идентификатор документа в БД",
        example=123
    )
    chunk_count: int = Field(
        0, 
        description="Количество проиндексированных чанков",
        example=15
    )
    error: Optional[str] = Field(
        None, 
        description="Сообщение об ошибке, если есть"
    )


class VisionResponse(BaseModel):
    """Модель ответа на распознавание изображения."""
    text: str = Field(
        ..., 
        description="Распознанный текст",
        example="Текст с изображения..."
    )
    status: str = Field(
        ..., 
        description="Статус: ok, no_text_found, error",
        example="ok"
    )
    error: Optional[str] = Field(
        None, 
        description="Сообщение об ошибке"
    )


class TranscribeResponse(BaseModel):
    """Модель ответа на транскрибацию аудио."""
    text: str = Field(
        ..., 
        description="Распознанный текст",
        example="Распознанная речь..."
    )
    status: str = Field(
        ..., 
        description="Статус: ok, no_speech_detected, error",
        example="ok"
    )
    error: Optional[str] = Field(
        None, 
        description="Сообщение об ошибке"
    )


# ============================================================
# ГЛОБАЛЬНЫЙ ОРКЕСТРАТОР (инициализируется один раз при старте)
# ============================================================

_orchestrator: Optional[RAGOrchestrator] = None
_config = None


def get_orchestrator() -> RAGOrchestrator:
    """Возвращает экземпляр RAGOrchestrator (инициализирует при первом вызове)."""
    global _orchestrator
    if _orchestrator is None:
        raise RuntimeError("RAGOrchestrator не инициализирован. Сервер не запущен корректно.")
    return _orchestrator


def init_orchestrator(config):
    """Инициализирует все сервисы и создаёт RAGOrchestrator."""
    global _orchestrator, _config
    _config = config
    logger.info("Инициализация RAG-сервера...")

    # ---- 1. Инициализация сервисов ----
    # PostgreSQL (история, документы)
    db = PostgresService(
        host=config.db_host,
        port=config.db_port,
        user=config.db_user,
        password=config.db_password,
        db_name=config.db_name
    )

    # Qdrant (векторный поиск)
    qdrant = QdrantService(
        host=config.qdrant_host,
        port=config.qdrant_port,
        grpc_port=config.qdrant_grpc_port,
        collection_name=config.qdrant_collection,
        vector_size=config.qdrant_vector_size,
        distance=config.qdrant_distance,
        prefer_grpc=False
    )

    # Эмбеддинги (GigaChat)
    embedding = EmbeddingService(
        api_key=config.gigachat_api_key,
        model=config.embedding_model,
        timeout=config.embedding_timeout,
        cache_size=config.embedding_cache_size,
        verify_ssl=config.embedding_verify_ssl
    )

    # База знаний (индексация, поиск)
    kb = KBService(
        db=db,
        qdrant=qdrant,
        embedding=embedding,
        collection_name=config.qdrant_collection,
        chunk_size_tokens=config.chunk_size_tokens,
        overlap_tokens=config.overlap_tokens,
        approx_chunk_chars=config.chunking_approx_chunk_chars,
        approx_overlap_chars=config.chunking_approx_overlap_chars
    )

    # GigaChat клиент (генерация)
    giga = GigaClient(
        api_key=config.gigachat_api_key,
        model=config.ai_model,
        temperature=config.ai_temperature,
        timeout=config.ai_timeout,
        verify_ssl=False
    )

    # Файловый сервис (извлечение текста)
    files = FileService(config)

    # ---- 2. Создание оркестратора ----
    # Загружаем дефолтные промпты (опционально)
    default_prompts = {}
    system_prompt_path = getattr(config, 'system_prompt_file', None)
    if system_prompt_path and system_prompt_path.exists():
        try:
            with open(system_prompt_path, 'r', encoding='utf-8') as f:
                default_prompts['system'] = f.read()
        except Exception as e:
            logger.warning(f"Не удалось загрузить системный промпт: {e}")

    orchestrator = RAGOrchestrator(
        db=db,
        qdrant=qdrant,
        embedding=embedding,
        kb=kb,
        giga=giga,
        files=files,
        config=config,
        default_prompts=default_prompts
    )

    _orchestrator = orchestrator
    logger.info("RAG-сервер успешно инициализирован")
    return orchestrator


# ============================================================
# FASTAPI ПРИЛОЖЕНИЕ С КАСТОМНОЙ ОPENAPI
# ============================================================

# Создаём приложение с базовыми метаданными
app = FastAPI(
    title="Эфцекабот RAG-сервер",
    description="""
    ## Единый RAG-сервис для интеллектуальной обработки запросов
    
    **Эфцекабот** – корпоративная платформа для гибридного RAG на базе GigaChat.
    
    ### Основные возможности
    - **Гибридный поиск** (dense + sparse) с RRF-ранжированием
    - **Самокритика и перегенерация** ответов
    - **Иерархическое резюмирование** длинных документов и диалогов
    - **Поддержка личных, комнатных и глобальных документов**
    - **Распознавание текста на изображениях (OCR)**
    - **Транскрибация аудио** через SaluteSpeech
    
    ### Аутентификация
    Все эндпоинты (кроме `/health`) требуют передачи API-ключа в заголовке `X-API-Key`.
    """,
    version="3.0.0",
    contact={
        "name": "Команда Эфцекабот",
        "email": "support@fckbot.ru",
    },
    license_info={
        "name": "ISC License",
        "url": "https://opensource.org/licenses/ISC",
    },
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_url="/openapi.json",
)

# Разрешаем CORS для возможности запросов с других доменов (для тестов)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Кастомизируем OpenAPI-схему для добавления примеров и кодов ошибок
def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    
    openapi_schema = get_openapi(
        title=app.title,
        version=app.version,
        description=app.description,
        routes=app.routes,
        contact=app.contact,
        license_info=app.license_info,
    )
    
    # Добавляем глобальные компоненты безопасности
    openapi_schema["components"]["securitySchemes"] = {
        "APIKeyHeader": {
            "type": "apiKey",
            "in": "header",
            "name": "X-API-Key",
            "description": "API-ключ для доступа к сервису. Получите у администратора.",
        }
    }
    
    # Добавляем глобальную security requirement
    openapi_schema["security"] = [{"APIKeyHeader": []}]
    
    # Исключаем /health из требований безопасности
    # (это делается через параметр `include_in_schema=False` в эндпоинте)
    
    app.openapi_schema = openapi_schema
    return app.openapi_schema


app.openapi = custom_openapi


# ============================================================
# LIFECYCLE EVENTS
# ============================================================

@app.on_event("startup")
async def startup_event():
    """Подключаемся к БД при старте сервера."""
    if _orchestrator is None:
        logger.error("Оркестратор не инициализирован.")
    else:
        try:
            await _orchestrator.db.connect()
            logger.info("Подключение к PostgreSQL установлено")
        except Exception as e:
            logger.error(f"Ошибка подключения к PostgreSQL: {e}")


@app.on_event("shutdown")
async def shutdown_event():
    """Закрываем соединение с БД при завершении сервера."""
    if _orchestrator:
        try:
            await _orchestrator.db.close()
            logger.info("Подключение к PostgreSQL закрыто")
        except Exception as e:
            logger.error(f"Ошибка закрытия PostgreSQL: {e}")


# ============================================================
# ЭНДПОИНТЫ API С ТЕГАМИ И ПОДРОБНЫМИ ОПИСАНИЯМИ
# ============================================================

# Теги для группировки в документации
tags_metadata = [
    {
        "name": "RAG",
        "description": "Основные RAG-операции: запросы и индексация документов",
    },
    {
        "name": "Vision",
        "description": "Распознавание текста на изображениях (OCR)",
    },
    {
        "name": "Audio",
        "description": "Транскрибация аудиофайлов (SaluteSpeech)",
    },
    {
        "name": "Health",
        "description": "Проверка состояния сервиса",
    },
]


@app.post(
    "/rag/query",
    response_model=QueryResponse,
    tags=["RAG"],
    summary="Выполнить RAG-запрос",
    description="""
    Отправляет текстовый запрос пользователя, выполняет полный RAG-пайплайн:
    
    1. **Классификация намерений** – определяет тип запроса (FACT, PROCEDURE, METRICS и т.д.)
    2. **Расширение запроса** – добавляет синонимы и связанные термины
    3. **Гибридный поиск** (dense + sparse) в Qdrant
    4. **Переранжирование** через кросс-энкодер
    5. **Генерация ответа** через GigaChat
    6. **Самокритика** и перегенерация при необходимости
    7. **Сохранение истории** диалога в PostgreSQL
    
    **Требования**: заголовок `X-API-Key` с валидным ключом.
    **История**: сервер автоматически получает историю диалога из БД.
    """,
    response_description="Ответ содержит сгенерированный текст, намерение, источники и уверенность"
)
async def rag_query(request: QueryRequest, _: str = Depends(verify_api_key)):
    """
    Выполняет RAG-запрос.
    """
    orchestrator = get_orchestrator()

    try:
        result = await orchestrator.process_query(
            query=request.query,
            user_jid=request.user_jid,
            room_jid=request.room_jid,
            prompts=request.prompts,
            intent_override=request.intent_override,
            last_file_path=request.last_file_path,
            last_file_text=request.last_file_text
        )
        return QueryResponse(
            answer=result.get('answer', ''),
            intent=result.get('intent', 'GENERAL'),
            context=result.get('context'),
            sources=result.get('sources', []),
            confidence=result.get('confidence'),
            error=result.get('error')
        )
    except Exception as e:
        logger.error(f"Ошибка в /rag/query: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=str(e)
        )


@app.post(
    "/rag/index",
    response_model=IndexResponse,
    tags=["RAG"],
    summary="Индексировать документ",
    description="""
    Индексирует документ в базу знаний.
    
    **Процесс:**
    1. Документ разбивается на чанки (с учётом токенов)
    2. Для каждого чанка вычисляются dense- и sparse-эмбеддинги
    3. Чанки сохраняются в Qdrant с метаданными
    4. Метаданные сохраняются в PostgreSQL с флагом `indexed`
    
    **Поддерживаемые типы:**
    - Личные (`is_global=False`, `room_jid=None`)
    - Комнатные (`room_jid` указан)
    - Глобальные (`is_global=True`, только для админов)
    
    **Обновление:** если документ с таким именем уже существует в комнате,
    он будет заменён новой версией (по хешу).
    """,
    response_description="Возвращает идентификатор документа и количество чанков"
)
async def index_document(request: IndexRequest, _: str = Depends(verify_api_key)):
    """Индексирует документ в базу знаний."""
    orchestrator = get_orchestrator()

    try:
        result = await orchestrator.index_document(
            file_name=request.file_name,
            file_text=request.file_text,
            user_jid=request.user_jid,
            room_jid=request.room_jid,
            is_global=request.is_global,
            title=request.title,
            metadata=request.metadata,
            file_hash=request.file_hash,
            update_if_exists=request.update_if_exists
        )
        return IndexResponse(
            doc_id=result.get('doc_id'),
            chunk_count=result.get('chunk_count', 0),
            error=result.get('error')
        )
    except Exception as e:
        logger.error(f"Ошибка в /rag/index: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=str(e)
        )


@app.post(
    "/rag/vision",
    response_model=VisionResponse,
    tags=["Vision"],
    summary="Распознать текст на изображении",
    description="""
    Отправляет изображение на OCR-распознавание через Tesseract.
    
    **Поддерживаемые форматы:**
    - JPEG, PNG, BMP, TIFF
    
    **Возвращает:**
    - Распознанный текст
    - Статус: `ok`, `no_text_found`, `error`
    
    **Требования:** заголовок `X-API-Key`.
    """,
    response_description="Текст, распознанный на изображении"
)
async def vision_endpoint(
    file: UploadFile = FastAPIFile(
        ..., 
        description="Файл изображения (JPEG, PNG, BMP, TIFF)"
    ),
    _: str = Depends(verify_api_key)
):
    orchestrator = get_orchestrator()
    
    # Проверяем расширение
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff'):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Поддерживаются только изображения (jpg, png, bmp, tiff)"
        )
    
    # Сохраняем во временный файл
    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
        tmp_path = tmp.name
        shutil.copyfileobj(file.file, tmp)
    
    try:
        loop = asyncio.get_running_loop()
        text = await loop.run_in_executor(
            None,
            orchestrator.files.extract_text_from_image,
            tmp_path
        )
        if not text:
            return VisionResponse(text="", status="no_text_found")
        return VisionResponse(text=text, status="ok")
    except Exception as e:
        logger.error(f"Ошибка OCR: {e}", exc_info=True)
        return VisionResponse(text="", status="error", error=str(e))
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


@app.post(
    "/rag/transcribe",
    response_model=TranscribeResponse,
    tags=["Audio"],
    summary="Транскрибировать аудио",
    description="""
    Отправляет аудиофайл на транскрибацию через SaluteSpeech.
    
    **Поддерживаемые форматы:**
    - OGG, WAV, MP3, AMR, M4A
    
    **Требования:**
    - Заголовок `X-API-Key`
    - Переменная окружения `SALUTE_SPEECH_AUTH`
    - Установленный `ffmpeg` на сервере
    
    **Возвращает:**
    - Распознанный текст
    - Статус: `ok`, `no_speech_detected`, `error`
    """,
    response_description="Текст, распознанный из аудио"
)
async def transcribe_endpoint(
    file: UploadFile = FastAPIFile(
        ..., 
        description="Аудиофайл (OGG, WAV, MP3, AMR, M4A)"
    ),
    _: str = Depends(verify_api_key)
):
    orchestrator = get_orchestrator()
    
    # Проверяем расширение
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in ('.ogg', '.wav', '.mp3', '.amr', '.m4a'):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Поддерживаются только аудио (ogg, wav, mp3, amr, m4a)"
        )
    
    # Сохраняем во временный файл
    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
        tmp_path = tmp.name
        shutil.copyfileobj(file.file, tmp)
    
    try:
        text = await orchestrator.files.transcribe_audio(tmp_path)
        if not text:
            return TranscribeResponse(text="", status="no_speech_detected")
        return TranscribeResponse(text=text, status="ok")
    except Exception as e:
        logger.error(f"Ошибка транскрибации: {e}", exc_info=True)
        return TranscribeResponse(text="", status="error", error=str(e))
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


@app.get(
    "/health",
    tags=["Health"],
    summary="Проверка состояния сервиса",
    description="""
    Проверяет, запущен ли RAG-сервер и инициализирован ли оркестратор.
    
    **Ответ:**
    - `healthy` – сервер работает
    - `unhealthy` – оркестратор не инициализирован
    
    **Примечание:** этот эндпоинт **не требует** аутентификации.
    """,
    response_description="Статус сервиса"
)
async def health_check():
    """Проверка здоровья сервиса (без аутентификации)."""
    if _orchestrator is None:
        return {"status": "unhealthy", "message": "Orchestrator not initialized"}
    return {"status": "healthy"}


# ============================================================
# ТОЧКА ВХОДА
# ============================================================

def main():
    """Главная функция запуска RAG-сервера."""
    import argparse

    parser = argparse.ArgumentParser(description="RAG-сервер")
    parser.add_argument(
        "--profile-dir",
        required=True,
        help="Путь к директории профиля бота (например, /usr/local/etc/bots/metabot)"
    )
    parser.add_argument(
        "--host",
        default="0.0.0.0",
        help="Хост для привязки сервера"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8080,
        help="Порт для привязки сервера"
    )
    parser.add_argument(
        "--log-level",
        default="info",
        choices=["debug", "info", "warning", "error", "critical"],
        help="Уровень логирования"
    )
    args = parser.parse_args()

    # Загружаем конфигурацию через Pydantic
    try:
        config = load_config(args.profile_dir)
    except Exception as e:
        print(f"Ошибка загрузки конфигурации: {e}")
        sys.exit(1)

    # Настраиваем логирование
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    setup_logging(config.name, config.log_file_path, log_level=log_level)

    logger.info(f"Запуск RAG-сервера с профилем {args.profile_dir}")
    logger.info(f"Хост: {args.host}, порт: {args.port}")

    # Устанавливаем конфиг для аутентификации
    set_auth_config(config)

    # Инициализируем оркестратор
    try:
        init_orchestrator(config)
    except Exception as e:
        logger.error(f"Ошибка инициализации оркестратора: {e}", exc_info=True)
        sys.exit(1)

    # Запускаем Uvicorn
    uvicorn.run(
        "rag.rag_server:app",
        host=args.host,
        port=args.port,
        log_level=args.log_level,
        reload=False
    )


if __name__ == "__main__":
    main()