"""
知识库管理器 - 集合管理 Mixin

提供向量库（ChromaDB Collection）的创建、删除、查询等管理功能。
每个向量库对应一个独立的 ChromaDB 集合，支持按部门隔离。

主要方法：
- get_collection: 获取或创建向量库集合
- create_collection: 创建新向量库
- delete_collection: 删除向量库
- list_collections: 列出所有向量库
- collection_exists: 检查向量库是否存在
"""

import os
import gc
import time
import logging
from typing import List, Optional, Tuple
from datetime import datetime

import chromadb
from chromadb import Collection

from .base import (
    CollectionInfo, PUBLIC_KB_NAME
)

logger = logging.getLogger(__name__)


class CollectionMixin:
    """
    集合管理 Mixin

    提供 ChromaDB 向量库集合的管理功能，包括：
    - 集合的创建与删除
    - 集合元数据管理
    - 集合列表查询

    依赖属性（需由主类提供）：
    - self.base_path: 向量库存储根路径
    - self.bm25_base_path: BM25 索引存储路径
    - self._collections: 集合缓存字典
    - self._clients: 客户端缓存字典
    - self._bm25_indexes: BM25 索引缓存字典
    - self._metadata: 元数据字典
    - self._lock: 线程锁
    """

    def _get_client(self, kb_name: str) -> chromadb.PersistentClient:
        """
        获取或创建 ChromaDB 客户端

        每个向量库使用独立的客户端实例，客户端会被缓存以避免重复创建。

        Args:
            kb_name: 向量库名称

        Returns:
            ChromaDB PersistentClient 实例
        """
        if kb_name not in self._clients:
            db_path = os.path.join(self.base_path, kb_name)
            os.makedirs(db_path, exist_ok=True)
            self._clients[kb_name] = chromadb.PersistentClient(path=db_path)
        return self._clients[kb_name]

    def get_collection(self, kb_name: str) -> Optional[Collection]:
        """
        获取或创建向量库集合

        如果集合已存在于缓存中，直接返回缓存实例。
        否则从 ChromaDB 获取或创建新集合。

        Args:
            kb_name: 向量库名称（只能包含字母、数字和下划线）

        Returns:
            ChromaDB Collection 对象；失败时返回 None

        Note:
            使用 cosine 相似度作为向量距离度量。
            sync_threshold 设置为 100000 以支持大批量写入。
        """
        with self._lock:
            if kb_name in self._collections:
                return self._collections[kb_name]
            try:
                client = self._get_client(kb_name)
                collection = client.get_or_create_collection(
                    name=kb_name,
                    metadata={
                        "hnsw:space": "cosine",
                        "hnsw:sync_threshold": 100000
                    }
                )
                self._collections[kb_name] = collection
                logger.info(f"获取向量库: {kb_name}, 文档数: {collection.count()}")
                return collection
            except Exception as e:
                logger.error(f"获取向量库失败: {kb_name}, 错误: {e}")
                return None

    def create_collection(
        self,
        kb_name: str,
        display_name: str = "",
        department: str = "",
        description: str = ""
    ) -> Tuple[bool, str]:
        """
        创建新向量库

        创建一个新的 ChromaDB 集合，同时初始化对应的 BM25 索引。

        Args:
            kb_name: 向量库名称（只能包含字母、数字和下划线）
            display_name: 显示名称（用于 UI 展示）
            department: 所属部门（用于权限控制）
            description: 向量库描述

        Returns:
            元组 (success, message)：
            - success: 创建是否成功
            - message: 成功或失败的描述信息

        Example:
            >>> success, msg = kb_manager.create_collection(
            ...     "dept_finance",
            ...     display_name="财务部知识库",
            ...     department="finance"
            ... )
        """
        from .base import BM25Index

        if not kb_name or not kb_name.replace('_', '').isalnum():
            return False, "向量库名称只能包含字母、数字和下划线"

        if kb_name in self._metadata.get("collections", {}):
            return False, f"向量库 '{kb_name}' 已存在"

        chroma_dir = os.path.join(self.base_path, kb_name)
        if os.path.exists(chroma_dir):
            import shutil
            shutil.rmtree(chroma_dir)
            logger.warning(f"清理残留向量库文件夹: {chroma_dir}")

        try:
            collection = self.get_collection(kb_name)
            if not collection:
                return False, "创建向量库失败"

            self._bm25_indexes[kb_name] = BM25Index()

            if "collections" not in self._metadata:
                self._metadata["collections"] = {}

            self._metadata["collections"][kb_name] = {
                "display_name": display_name or kb_name,
                "department": department,
                "description": description,
                "created_at": datetime.now().isoformat()
            }
            self._save_metadata()

            logger.info(f"创建向量库: {kb_name}")
            return True, f"向量库 '{kb_name}' 创建成功"

        except Exception as e:
            logger.error(f"创建向量库失败: {e}")
            return False, f"创建失败: {str(e)}"

    def update_collection_metadata(
        self,
        kb_name: str,
        display_name: str = None,
        description: str = None
    ) -> bool:
        """
        更新向量库元数据

        仅更新 display_name 和 description，不影响集合中的数据。

        Args:
            kb_name: 向量库名称
            display_name: 新的显示名称（None 表示不更新）
            description: 新的描述（None 表示不更新）

        Returns:
            更新成功返回 True，向量库不存在返回 False
        """
        collections = self._metadata.get("collections", {})
        if kb_name not in collections:
            return False

        if display_name is not None:
            collections[kb_name]["display_name"] = display_name
        if description is not None:
            collections[kb_name]["description"] = description

        self._save_metadata()
        logger.info(f"更新向量库元数据: {kb_name}")
        return True

    def delete_collection(self, kb_name: str, delete_documents: bool = False) -> Tuple[bool, str]:
        """
        删除向量库

        删除向量库及其所有数据，包括：
        - ChromaDB 集合
        - BM25 索引文件
        - 向量库文件夹
        - 可选：原始文档文件

        Args:
            kb_name: 向量库名称
            delete_documents: 是否同时删除原始文档文件

        Returns:
            元组 (success, message)：
            - success: 删除是否成功
            - message: 成功或失败的描述信息

        Warning:
            公开知识库 (public_kb) 不能被删除。
            删除操作不可逆，请谨慎使用。
        """
        import shutil

        if kb_name == PUBLIC_KB_NAME:
            return False, "公开知识库不能删除"

        if kb_name not in self._metadata.get("collections", {}):
            return False, f"向量库 '{kb_name}' 不存在"

        try:
            client = self._clients.get(kb_name)
            if client:
                try:
                    client.delete_collection(kb_name)
                except Exception as e:
                    logger.debug(f"删除集合失败: {e}")
                try:
                    client.close()
                except Exception as e:
                    logger.debug(f"关闭客户端失败: {e}")

            if kb_name in self._collections:
                del self._collections[kb_name]
            if kb_name in self._bm25_indexes:
                del self._bm25_indexes[kb_name]
            if kb_name in self._clients:
                del self._clients[kb_name]

            gc.collect()
            time.sleep(1)

            bm25_path = os.path.join(self.bm25_base_path, f"{kb_name}.pkl")
            if os.path.exists(bm25_path):
                os.remove(bm25_path)

            chroma_dir = os.path.join(self.base_path, kb_name)
            if os.path.exists(chroma_dir):
                for attempt in range(3):
                    try:
                        shutil.rmtree(chroma_dir)
                        logger.info(f"删除向量库文件夹: {chroma_dir}")
                        break
                    except PermissionError:
                        if attempt < 2:
                            logger.warning(f"删除文件夹失败，重试 {attempt + 1}/3")
                            gc.collect()
                            time.sleep(1)
                        else:
                            logger.error(f"删除文件夹失败，文件被占用: {chroma_dir}")
                            raise

            if delete_documents:
                from config import DOCUMENTS_PATH
                docs_dir = os.path.join(DOCUMENTS_PATH, kb_name)
                if os.path.exists(docs_dir):
                    shutil.rmtree(docs_dir)
                    logger.info(f"删除文档文件夹: {docs_dir}")

            # 删除该向量库所有文档的哈希记录
            try:
                from knowledge.sync import SyncDatabase
                sync_db = SyncDatabase()
                # 获取所有以 kb_name/ 开头的哈希记录
                all_hashes = sync_db.get_all_document_hashes()
                deleted_count = 0
                for doc_id in list(all_hashes.keys()):
                    if doc_id.startswith(f"{kb_name}/"):
                        sync_db.delete_document_hash(doc_id)
                        deleted_count += 1
                if deleted_count > 0:
                    logger.info(f"清理向量库哈希记录: {kb_name}, 共 {deleted_count} 条")
            except Exception as e:
                logger.warning(f"清理哈希记录失败: {e}")

            # 清理该向量库的 SQLite 版本记录和变更日志
            try:
                from data.db import get_connection
                with get_connection("knowledge") as conn:
                    cursor1 = conn.execute(
                        "DELETE FROM document_versions WHERE collection = ?",
                        (kb_name,)
                    )
                    cursor2 = conn.execute(
                        "DELETE FROM version_change_logs WHERE collection = ?",
                        (kb_name,)
                    )
                    conn.commit()
                    ver_cleaned = cursor1.rowcount
                    log_cleaned = cursor2.rowcount
                    if ver_cleaned > 0 or log_cleaned > 0:
                        logger.info(
                            f"清理向量库版本记录: {kb_name}, "
                            f"版本 {ver_cleaned} 条, 日志 {log_cleaned} 条"
                        )
            except Exception as e:
                logger.warning(f"清理版本记录失败: {e}")

            if kb_name in self._metadata.get("collections", {}):
                del self._metadata["collections"][kb_name]
            self._save_metadata()

            logger.info(f"删除向量库: {kb_name}")
            return True, f"向量库 '{kb_name}' 已删除"

        except Exception as e:
            logger.error(f"删除向量库失败: {e}")
            return False, f"删除失败: {str(e)}"

    def list_collections(self) -> List[CollectionInfo]:
        """
        列出所有向量库

        返回所有已创建的向量库信息，包括名称、文档数量、创建时间等。

        Returns:
            CollectionInfo 对象列表，每个对象包含：
            - name: 向量库名称
            - display_name: 显示名称
            - document_count: 文档数量
            - created_at: 创建时间
            - department: 所属部门
            - description: 描述
        """
        result = []

        # 扫描 base_path 下的所有子目录作为向量库
        # 每个向量库使用独立目录：base_path/my_ky, base_path/public_kb 等
        actual_collections = []
        try:
            if os.path.exists(self.base_path):
                for item in os.listdir(self.base_path):
                    item_path = os.path.join(self.base_path, item)
                    if os.path.isdir(item_path) and not item.startswith('.'):
                        # 检查是否包含 chroma.sqlite3（有效的向量库目录）
                        if os.path.exists(os.path.join(item_path, 'chroma.sqlite3')):
                            actual_collections.append(item)
        except Exception as e:
            logger.warning(f"扫描向量库目录失败: {e}")

        # 如果扫描失败，回退到元数据中的集合列表
        if not actual_collections:
            actual_collections = list(self._metadata.get("collections", {}).keys())

        for name in actual_collections:
            if name not in self._metadata.get("collections", {}):
                if "collections" not in self._metadata:
                    self._metadata["collections"] = {}
                self._metadata["collections"][name] = {
                    "display_name": name,
                    "department": "",
                    "description": "",
                    "created_at": datetime.now().isoformat()
                }
                logger.info(f"自动补充向量库元数据: {name}")

        self._save_metadata()

        for name, info in self._metadata.get("collections", {}).items():
            collection = self.get_collection(name)
            result.append(CollectionInfo(
                name=name,
                display_name=info.get("display_name", name),
                document_count=collection.count() if collection else 0,
                created_at=info.get("created_at", ""),
                department=info.get("department", ""),
                description=info.get("description", "")
            ))

        return result

    def collection_exists(self, kb_name: str) -> bool:
        """
        检查向量库是否存在

        Args:
            kb_name: 向量库名称

        Returns:
            存在返回 True，不存在返回 False
        """
        return kb_name in self._metadata.get("collections", {})