# rag/services/postgres_service.py
"""
Сервис для работы с PostgreSQL (через asyncpg).
Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
управление доступом к личным документам.

Добавлена поддержка флага indexed для транзакционности индексации.
Добавлена колонка file_text для хранения исходного текста документа (для переиндексации).
"""

import asyncio
import json
import logging
from typing import List, Dict, Any, Optional

import asyncpg
from asyncpg import Pool, Connection

logger = logging.getLogger(__name__)


class PostgresService:
    """
    Асинхронный сервис для работы с PostgreSQL.
    Использует пул соединений asyncpg.
    """

    def __init__(self, host: str, port: int, user: str, password: str, db_name: str):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.db_name = db_name
        self.pool: Optional[Pool] = None

    # ------------------------------------------------------------------
    # Подключение и проверка таблиц
    # ------------------------------------------------------------------

    async def connect(self, create_if_missing: bool = False) -> None:
        """
        Создаёт пул соединений и проверяет наличие таблиц.
        Таблицы должны быть созданы через Alembic.

        Аргументы:
            create_if_missing: если True, создаёт таблицы автоматически (только для тестов).
                               По умолчанию False.
        """
        self.pool = await asyncpg.create_pool(
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            database=self.db_name,
            min_size=1,
            max_size=10,
            command_timeout=60
        )
        if create_if_missing:
            await self._init_tables()
            logger.warning("Таблицы созданы автоматически (режим тестирования). В продакшене используйте Alembic.")
        else:
            await self._check_tables()
        logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")

    async def _check_tables(self) -> None:
        """Проверяет наличие основных таблиц."""
        required_tables = ['documents', 'document_access', 'history', 'rooms', 'room_templates']
        missing = []
        async with self.pool.acquire() as conn:
            for table in required_tables:
                row = await conn.fetchval(
                    "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)",
                    table
                )
                if not row:
                    missing.append(table)
        if missing:
            logger.warning(
                f"Отсутствуют таблицы: {', '.join(missing)}. "
                "Запустите миграции Alembic: `alembic upgrade head`"
            )
        else:
            logger.debug("Все необходимые таблицы существуют")

    async def _init_tables(self) -> None:
        """Создаёт таблицы (только для тестов)."""
        async with self.pool.acquire() as conn:
            # 1. Таблица documents с колонкой file_text
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id SERIAL PRIMARY KEY,
                    source_name VARCHAR(255) NOT NULL,
                    upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    owner_jid VARCHAR(255) NOT NULL,
                    is_global BOOLEAN DEFAULT FALSE,
                    collection_name VARCHAR(64) NOT NULL,
                    metadata JSONB,
                    room_jid VARCHAR(255) DEFAULT NULL,
                    file_hash VARCHAR(64) DEFAULT NULL,
                    indexed BOOLEAN DEFAULT FALSE,
                    file_text TEXT DEFAULT NULL
                )
            """)
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner ON documents(owner_jid)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_global ON documents(is_global)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room ON documents(room_jid)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_indexed ON documents(indexed)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_room_source ON documents(room_jid, source_name)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_documents_owner_source ON documents(owner_jid, source_name)")
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS ft_source_name 
                ON documents USING GIN (to_tsvector('russian', source_name))
            """)

            # 2. Таблица document_access
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS document_access (
                    doc_id INT,
                    user_jid VARCHAR(255),
                    PRIMARY KEY (doc_id, user_jid)
                )
            """)
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_access_user ON document_access(user_jid)")

            # 3. Тип ENUM для history.role
            await conn.execute("""
                DO $$ BEGIN
                    CREATE TYPE history_role AS ENUM ('user', 'assistant');
                EXCEPTION
                    WHEN duplicate_object THEN NULL;
                END $$;
            """)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS history (
                    id SERIAL PRIMARY KEY,
                    user_jid VARCHAR(255) NOT NULL,
                    role history_role NOT NULL,
                    content TEXT,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    room_jid VARCHAR(255) DEFAULT NULL
                )
            """)
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_user ON history(user_jid)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_room ON history(room_jid)")
            await conn.execute("CREATE INDEX IF NOT EXISTS idx_history_time ON history(timestamp)")

            # 4. Таблица rooms
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS rooms (
                    room_jid VARCHAR(255) PRIMARY KEY
                )
            """)

            # 5. Таблица room_templates
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS room_templates (
                    room_jid VARCHAR(255) NOT NULL,
                    name VARCHAR(255) NOT NULL,
                    file_path TEXT NOT NULL,
                    file_hash VARCHAR(64),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    PRIMARY KEY (room_jid, name)
                )
            """)
            logger.info("Таблицы созданы (или уже существовали)")

    async def close(self) -> None:
        """Закрывает пул соединений."""
        if self.pool:
            await self.pool.close()

    # ------------------------------------------------------------------
    # Документы
    # ------------------------------------------------------------------

    async def add_document(
        self,
        source_name: str,
        owner_jid: str,
        is_global: bool,
        collection_name: str,
        metadata: dict = None,
        room_jid: str = None,
        file_hash: str = None,
        indexed: bool = False,
        file_text: Optional[str] = None
    ) -> int:
        """
        Добавляет запись о документе с возможностью сохранить исходный текст.
        Возвращает doc_id.
        """
        metadata_json = json.dumps(metadata, ensure_ascii=False) if metadata else None
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                INSERT INTO documents (
                    source_name, owner_jid, is_global, collection_name,
                    metadata, room_jid, file_hash, indexed, file_text
                )
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                RETURNING id
                """,
                source_name, owner_jid, is_global, collection_name,
                metadata_json, room_jid, file_hash, indexed, file_text
            )
            return row['id']

    async def update_document_indexed(self, doc_id: int, indexed: bool) -> None:
        """Обновляет флаг indexed для документа."""
        async with self.pool.acquire() as conn:
            await conn.execute(
                "UPDATE documents SET indexed = $1 WHERE id = $2",
                indexed, doc_id
            )

    async def get_document(self, doc_id: int) -> Optional[Dict]:
        """Возвращает запись документа по id."""
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow("SELECT * FROM documents WHERE id = $1", doc_id)
            return dict(row) if row else None

    async def add_access(self, doc_id: int, user_jid: str) -> None:
        """Даёт пользователю доступ к личному документу."""
        async with self.pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO document_access (doc_id, user_jid)
                VALUES ($1, $2)
                ON CONFLICT (doc_id, user_jid) DO NOTHING
                """,
                doc_id, user_jid
            )

    async def can_access(self, doc_id: int, user_jid: str, is_global: bool) -> bool:
        """Проверяет, имеет ли пользователь доступ к документу."""
        if is_global:
            return True
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT 1 FROM document_access WHERE doc_id = $1 AND user_jid = $2",
                doc_id, user_jid
            )
            return row is not None

    async def find_document_by_name_and_room(self, source_name: str, room_jid: str) -> Optional[Dict]:
        """
        Поиск документа в комнате по имени (source_name). Используется для обновления версий.
        Возвращает словарь с ключами id, file_hash или None.
        """
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT id, file_hash FROM documents WHERE source_name = $1 AND room_jid = $2 LIMIT 1",
                source_name, room_jid
            )
            return dict(row) if row else None

    # ------------------------------------------------------------------
    # История диалогов
    # ------------------------------------------------------------------

    async def add_history(self, user_jid: str, role: str, content: str, room_jid: str = None) -> None:
        """Сохраняет сообщение в историю."""
        async with self.pool.acquire() as conn:
            await conn.execute(
                "INSERT INTO history (user_jid, role, content, room_jid) VALUES ($1, $2, $3, $4)",
                user_jid, role, content, room_jid
            )

    async def get_history(self, user_jid: str, room_jid: str = None, limit: int = 100, offset: int = 0) -> List[Dict]:
        """
        Возвращает историю диалога с поддержкой пагинации.

        Аргументы:
            user_jid: JID пользователя.
            room_jid: JID комнаты (None для личного чата).
            limit: максимальное количество сообщений.
            offset: смещение (сколько пропустить).

        Возвращает:
            List[Dict]: список сообщений с ключами role, content, timestamp.
        """
        async with self.pool.acquire() as conn:
            if room_jid:
                rows = await conn.fetch(
                    """
                    SELECT role, content, timestamp 
                    FROM history 
                    WHERE room_jid = $1 
                    ORDER BY timestamp ASC 
                    LIMIT $2 OFFSET $3
                    """,
                    room_jid, limit, offset
                )
            else:
                rows = await conn.fetch(
                    """
                    SELECT role, content, timestamp 
                    FROM history 
                    WHERE user_jid = $1 AND room_jid IS NULL 
                    ORDER BY timestamp ASC 
                    LIMIT $2 OFFSET $3
                    """,
                    user_jid, limit, offset
                )
            return [dict(row) for row in rows]

    async def clear_user_history(self, user_jid: str) -> None:
        """Удаляет личную историю пользователя."""
        async with self.pool.acquire() as conn:
            await conn.execute("DELETE FROM history WHERE user_jid = $1 AND room_jid IS NULL", user_jid)

    async def clear_room_history(self, room_jid: str) -> None:
        """Удаляет историю указанной комнаты."""
        async with self.pool.acquire() as conn:
            await conn.execute("DELETE FROM history WHERE room_jid = $1", room_jid)

    # ------------------------------------------------------------------
    # Полнотекстовый поиск (fallback)
    # ------------------------------------------------------------------

    async def fulltext_search(self, query: str, user_jid: str, room_jid: str = None) -> List[Dict]:
        """
        Поиск по source_name с использованием полнотекстового индекса PostgreSQL.
        Возвращает список документов (id, source_name, metadata) только с indexed = TRUE.
        """
        if not query or len(query.strip()) < 3:
            return []
        async with self.pool.acquire() as conn:
            if room_jid:
                rows = await conn.fetch(
                    """
                    SELECT id, source_name, metadata
                    FROM documents
                    WHERE room_jid = $1
                      AND indexed = TRUE
                      AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
                    LIMIT 10
                    """,
                    room_jid, query
                )
            else:
                rows = await conn.fetch(
                    """
                    SELECT id, source_name, metadata
                    FROM documents
                    WHERE (owner_jid = $1 OR is_global = TRUE)
                      AND room_jid IS NULL
                      AND indexed = TRUE
                      AND to_tsvector('russian', source_name) @@ plainto_tsquery('russian', $2)
                    LIMIT 10
                    """,
                    user_jid, query
                )
            return [dict(row) for row in rows]

    # ------------------------------------------------------------------
    # Комнаты MUC
    # ------------------------------------------------------------------

    async def add_room(self, room_jid: str) -> None:
        """Добавляет комнату в список."""
        async with self.pool.acquire() as conn:
            await conn.execute(
                "INSERT INTO rooms (room_jid) VALUES ($1) ON CONFLICT (room_jid) DO NOTHING",
                room_jid
            )

    async def get_rooms(self) -> List[str]:
        """Возвращает список JID комнат, в которых бот должен состоять."""
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("SELECT room_jid FROM rooms")
            return [row['room_jid'] for row in rows]

    async def remove_room(self, room_jid: str) -> None:
        """Удаляет комнату из списка."""
        async with self.pool.acquire() as conn:
            await conn.execute("DELETE FROM rooms WHERE room_jid = $1", room_jid)

    # ------------------------------------------------------------------
    # Шаблоны
    # ------------------------------------------------------------------

    async def save_template(self, room_jid: str, name: str, file_path: str, file_hash: str) -> bool:
        """Сохраняет шаблон в БД."""
        async with self.pool.acquire() as conn:
            await conn.execute("""
                INSERT INTO room_templates (room_jid, name, file_path, file_hash)
                VALUES ($1, $2, $3, $4)
                ON CONFLICT (room_jid, name) DO UPDATE SET
                    file_path = EXCLUDED.file_path,
                    file_hash = EXCLUDED.file_hash,
                    created_at = CURRENT_TIMESTAMP
            """, room_jid, name, file_path, file_hash)
            return True

    async def get_template(self, room_jid: str, name: str) -> Optional[Dict]:
        """Возвращает словарь {file_path, file_hash} или None."""
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT file_path, file_hash FROM room_templates WHERE room_jid = $1 AND name = $2",
                room_jid, name
            )
            return dict(row) if row else None

    async def list_templates(self, room_jid: str) -> List[str]:
        """Возвращает список названий шаблонов в комнате."""
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("SELECT name FROM room_templates WHERE room_jid = $1 ORDER BY name", room_jid)
            return [row['name'] for row in rows]

    async def delete_template(self, room_jid: str, name: str) -> bool:
        """Удаляет шаблон."""
        async with self.pool.acquire() as conn:
            result = await conn.execute(
                "DELETE FROM room_templates WHERE room_jid = $1 AND name = $2",
                room_jid, name
            )
            return result == "DELETE 1"

    # ------------------------------------------------------------------
    # Получение списков документов для пользователей (только indexed = TRUE)
    # ------------------------------------------------------------------

    async def get_available_sources(self, user_jid: str, room_jid: str = None) -> List[str]:
        """Возвращает список названий документов (source_name), доступных пользователю."""
        async with self.pool.acquire() as conn:
            if room_jid:
                rows = await conn.fetch(
                    "SELECT source_name FROM documents WHERE room_jid = $1 AND indexed = TRUE ORDER BY source_name",
                    room_jid
                )
            else:
                rows = await conn.fetch(
                    "SELECT source_name FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL AND indexed = TRUE ORDER BY source_name",
                    user_jid
                )
            return [row['source_name'] for row in rows]

    async def get_sources_with_type(self, user_jid: str, room_jid: str = None) -> List[Dict[str, Any]]:
        """Возвращает список документов с указанием, является ли документ глобальным."""
        async with self.pool.acquire() as conn:
            if room_jid:
                rows = await conn.fetch(
                    "SELECT source_name, is_global FROM documents WHERE room_jid = $1 AND indexed = TRUE ORDER BY source_name",
                    room_jid
                )
            else:
                rows = await conn.fetch(
                    "SELECT source_name, is_global FROM documents WHERE (owner_jid = $1 OR is_global = TRUE) AND room_jid IS NULL AND indexed = TRUE ORDER BY source_name",
                    user_jid
                )
            return [{"source_name": row['source_name'], "is_global": row['is_global']} for row in rows]

    # ------------------------------------------------------------------
    # Очистка баз знаний (удаление всех документов, независимо от indexed)
    # ------------------------------------------------------------------

    async def clear_user_kb(self, user_jid: str) -> List[int]:
        """Удаляет все личные документы пользователя."""
        async with self.pool.acquire() as conn:
            rows = await conn.fetch(
                "SELECT id FROM documents WHERE owner_jid = $1 AND is_global = FALSE AND room_jid IS NULL",
                user_jid
            )
            doc_ids = [row['id'] for row in rows]
            if doc_ids:
                await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
            return doc_ids

    async def clear_room_kb(self, room_jid: str) -> List[int]:
        """Удаляет все документы указанной комнаты."""
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("SELECT id FROM documents WHERE room_jid = $1", room_jid)
            doc_ids = [row['id'] for row in rows]
            if doc_ids:
                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
            return doc_ids

    async def clear_global_kb(self) -> List[int]:
        """Удаляет все глобальные документы."""
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE")
            doc_ids = [row['id'] for row in rows]
            if doc_ids:
                await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)
            return doc_ids

    async def delete_documents_by_ids(self, doc_ids: List[int]) -> None:
        """Удаляет документы по списку id."""
        if not doc_ids:
            return
        async with self.pool.acquire() as conn:
            await conn.execute("DELETE FROM document_access WHERE doc_id = ANY($1)", doc_ids)
            await conn.execute("DELETE FROM documents WHERE id = ANY($1)", doc_ids)