#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RAG-сервер – отдельный HTTP-сервис для обработки RAG-запросов.
Предоставляет API для выполнения RAG-запросов и индексации документов.

Этот сервер запускается как отдельный процесс и может обслуживать
несколько ботов и других систем одновременно.

ИСПОЛЬЗУЕТ: FastAPI для создания REST API и Uvicorn для запуска.
ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL).

ДОБАВЛЕНО:
- Эндпоинт /rag/vision для распознавания текста на изображениях (OCR)
- Эндпоинт /rag/transcribe для транскрибации аудио (SaluteSpeech)
"""

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

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

from fastapi import FastAPI, HTTPException, Depends, UploadFile, File as FastAPIFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# Импортируем наши модули
from core.utils.config_loader import BotConfig
from core.utils.logger import setup_logging
from core.services.postgres_service import PostgresService
from core.services.qdrant_service import QdrantService
from core.services.embedding_service import EmbeddingService
from core.services.kb_service import KBService
from core.services.giga_client import GigaClient
from core.services.file_service import FileService
from core.rag_orchestrator import RAGOrchestrator

logger = logging.getLogger(__name__)


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

class QueryRequest(BaseModel):
    """
    Модель запроса к RAG-серверу на выполнение RAG-запроса.
    ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
    """
    query: str = Field(..., description="Текст запроса пользователя")
    user_jid: str = Field(..., description="JID пользователя (без ресурса)")
    room_jid: Optional[str] = Field(None, description="JID комнаты (None для личного чата)")

    # history: List[Dict[str, str]] = Field(...)  # УБРАНО!

    prompts: Dict[str, str] = Field(
        default_factory=dict,
        description="Словарь промптов (system, synthesis, intent, expand, critique и т.д.)"
    )

    intent_override: Optional[str] = Field(None, description="Принудительное переопределение намерения")
    last_file_path: Optional[str] = Field(None, description="Путь к последнему загруженному файлу")
    last_file_text: Optional[str] = Field(None, description="Текст последнего загруженного файла")


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


class IndexRequest(BaseModel):
    """Модель запроса на индексацию документа."""
    file_name: str = Field(..., description="Исходное имя файла")
    file_text: str = Field(..., description="Извлечённый текст документа")
    user_jid: str = Field(..., description="JID владельца документа")
    room_jid: Optional[str] = Field(None, description="JID комнаты (если документ комнатный)")
    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="Идентификатор документа в БД")
    chunk_count: int = Field(0, description="Количество проиндексированных чанков")
    error: Optional[str] = Field(None, description="Сообщение об ошибке, если есть")


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


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


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

_orchestrator: Optional[RAGOrchestrator] = None
_config: Optional[BotConfig] = None


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


def init_orchestrator(config: BotConfig) -> RAGOrchestrator:
    """Инициализирует все сервисы и создаёт 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 ПРИЛОЖЕНИЕ
# ============================================================

app = FastAPI(
    title="RAG-сервер",
    description="Единый сервис для RAG-обработки запросов и индексации документов",
    version="1.0.0"
)

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


@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
# ============================================================

@app.post("/rag/query", response_model=QueryResponse)
async def rag_query(request: QueryRequest):
    """
    Выполняет 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=500, detail=str(e))


@app.post("/rag/index", response_model=IndexResponse)
async def index_document(request: IndexRequest):
    """Индексирует документ в базу знаний."""
    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=500, detail=str(e))


@app.get("/health")
async def health_check():
    """Проверка здоровья сервиса."""
    if _orchestrator is None:
        return {"status": "unhealthy", "message": "Orchestrator not initialized"}
    return {"status": "healthy"}


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

@app.post("/rag/vision", response_model=VisionResponse)
async def vision_endpoint(
    file: UploadFile = FastAPIFile(...),
    user_jid: str = None,
    room_jid: str = None
):
    """
    Распознаёт текст на изображении через OCR (Tesseract).
    Возвращает JSON: {"text": str, "status": "ok" | "no_text_found" | "error"}
    """
    orchestrator = get_orchestrator()
    
    # Проверяем расширение
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff'):
        raise HTTPException(status_code=400, detail="Поддерживаются только изображения (jpg, png, bmp, tiff)")
    
    # Сохраняем во временный файл
    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
        tmp_path = tmp.name
        shutil.copyfileobj(file.file, tmp)
    
    try:
        # Извлекаем текст через существующий FileService (синхронный метод)
        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)
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


@app.post("/rag/transcribe", response_model=TranscribeResponse)
async def transcribe_endpoint(
    file: UploadFile = FastAPIFile(...),
    user_jid: str = None,
    room_jid: str = None
):
    """
    Транскрибирует аудиофайл через SaluteSpeech.
    Возвращает JSON: {"text": str, "status": "ok" | "no_speech_detected" | "error"}
    """
    orchestrator = get_orchestrator()
    
    # Проверяем расширение
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in ('.ogg', '.wav', '.mp3', '.amr', '.m4a'):
        raise HTTPException(status_code=400, 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:
        # Транскрибируем через существующий FileService (асинхронный метод)
        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)
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


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

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()

    # Загружаем конфигурацию
    try:
        config = BotConfig(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, log_level=log_level)

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

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

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


if __name__ == "__main__":
    main()