# -*- coding: utf-8 -*-
"""
Сервис для работы с PostgreSQL (через asyncpg).
Обеспечивает пул соединений, CRUD операции для документов, истории диалогов, комнат MUC,
управление доступом к личным документам.

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

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) -> None:
        """
        Создаёт пул соединений и инициализирует таблицы, если их нет.
        """
        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
        )
        await self._init_tables()
        logger.info(f"Подключение к PostgreSQL БД {self.db_name} установлено (пул)")

    async def _init_tables(self) -> None:
        """
        Создаёт таблицы (если не существуют) с учётом всех необходимых колонок и индексов.
        Использует схему из schema_postgresql.sql (адаптированную для PostgreSQL).
        """
        async with self.pool.acquire() as conn:
            # 1. Таблица documents
            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
                )
            """)
            # Индексы
            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)")
            # Полнотекстовый индекс (GIN) – используем 'russian', при необходимости заменить на 'simple'
            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 $$;
            """)
            # Таблица history
            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
                )
            """)

    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   # новый параметр
    ) -> int:
        """
        Добавляет запись о документе. Возвращает doc_id.
        :param indexed: начальное состояние индексации (по умолчанию False)
        """
        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)
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                RETURNING id
                """,
                source_name, owner_jid, is_global, collection_name, metadata_json, room_jid, file_hash, indexed
            )
            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:
        """
        Сохраняет сообщение в историю.
        :param role: 'user' или 'assistant'
        """
        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) -> List[Dict]:
        """
        Возвращает последние limit сообщений, отсортированных от старых к новым.
        Если room_jid указан – история комнаты, иначе – личная история пользователя.
        """
        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
                    """,
                    room_jid, limit
                )
            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
                    """,
                    user_jid, limit
                )
            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, если Qdrant недоступен)
    # ------------------------------------------------------------------

    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:
        """Удаляет шаблон. Возвращает True, если удалён."""
        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), доступных пользователю.
        Только те, у которых indexed = TRUE.
        """
        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]]:
        """
        Возвращает список документов с указанием, является ли документ глобальным.
        Только indexed = TRUE.
        """
        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]:
        """Удаляет все личные документы пользователя (не глобальные и не комнатные). Возвращает список удалённых doc_id."""
        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]:
        """Удаляет все документы указанной комнаты. Возвращает список doc_id."""
        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]:
        """Удаляет все глобальные документы. Возвращает список doc_id."""
        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 (используется в global_remove)."""
        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)