"""
知识库管理器 - BM25索引管理 Mixin

提供 BM25 关键词检索索引的管理功能，包括索引的加载、保存和重建。

BM25 是一种基于概率检索模型的排序函数，适合中文关键词检索。
每个向量库对应一个独立的 BM25 索引，使用 jieba 进行中文分词。

主要方法：
- get_bm25_index: 获取或加载 BM25 索引
- save_bm25_index: 保存 BM25 索引到磁盘
- rebuild_bm25_index: 从向量库重建 BM25 索引
"""

import os
import logging
from typing import Optional

from .base import BM25Index

logger = logging.getLogger(__name__)


class IndexMixin:
    """
    BM25索引管理 Mixin

    提供 BM25 索引的生命周期管理，支持：
    - 索引的懒加载（首次访问时从磁盘加载）
    - 索引的持久化（保存到 .pkl 文件）
    - 索引的重建（从向量库同步）

    依赖属性（需由主类提供）：
    - self.bm25_base_path: BM25 索引存储路径
    - self._bm25_indexes: BM25 索引缓存字典
    """

    def get_bm25_index(self, kb_name: str) -> BM25Index:
        """
        获取或加载 BM25 索引

        如果索引已在内存中，直接返回。
        否则尝试从磁盘加载，加载失败则创建空索引。

        Args:
            kb_name: 向量库名称

        Returns:
            BM25Index 实例

        Note:
            索引文件路径: {bm25_base_path}/{kb_name}.pkl
        """
        if kb_name not in self._bm25_indexes:
            self._bm25_indexes[kb_name] = BM25Index()
            bm25_path = os.path.join(self.bm25_base_path, f"{kb_name}.pkl")
            if os.path.exists(bm25_path):
                self._bm25_indexes[kb_name].load(bm25_path)
        return self._bm25_indexes[kb_name]

    def save_bm25_index(self, kb_name: str) -> None:
        """
        保存 BM25 索引到磁盘

        将指定向量库的 BM25 索引序列化保存到 .pkl 文件。

        Args:
            kb_name: 向量库名称
        """
        if kb_name in self._bm25_indexes:
            bm25_path = os.path.join(self.bm25_base_path, f"{kb_name}.pkl")
            self._bm25_indexes[kb_name].save(bm25_path)
            logger.info(f"保存 BM25 索引: {kb_name}")

    def rebuild_bm25_index(self, kb_name: str) -> bool:
        """
        重建 BM25 索引

        从 ChromaDB 向量库中读取所有文档，重新构建 BM25 索引。
        用于文档变更后同步索引状态。

        Args:
            kb_name: 向量库名称

        Returns:
            重建成功返回 True，失败返回 False

        Warning:
            大量文档时重建可能耗时较长。
        """
        try:
            collection = self.get_collection(kb_name)
            if not collection:
                return False

            result = collection.get()

            bm25_index = BM25Index()
            if result['ids']:
                bm25_index.add_documents(
                    ids=result['ids'],
                    documents=result['documents'],
                    metadatas=result['metadatas']
                )

            self._bm25_indexes[kb_name] = bm25_index
            self.save_bm25_index(kb_name)

            logger.info(f"重建 BM25 索引: {kb_name}, 文档数: {len(result['ids'])}")
            return True

        except Exception as e:
            logger.error(f"重建 BM25 索引失败: {e}")
            return False

    def _rebuild_bm25_index(self, kb_name: str) -> None:
        """
        重建指定向量库的 BM25 索引（内部方法）

        兼容旧代码的内部方法，直接调用 rebuild_bm25_index。

        Args:
            kb_name: 向量库名称
        """
        self.rebuild_bm25_index(kb_name)