"""
BM25存储管理模块
使用单例模式确保全局共享同一个BM25检索器实例
"""
from typing import Optional, List
from pathlib import Path
from loguru import logger
from llama_index.core.storage.docstore import SimpleDocumentStore

# 使用自定义的中文 BM25 检索器
try:
    from src.retrievers import ChineseBM25Retriever
    HAS_BM25_RETRIEVER = True
except ImportError:
    HAS_BM25_RETRIEVER = False
    ChineseBM25Retriever = None

from src.config import get_config_manager, ConfigManager
from src.tokenizers import create_tokenizer
from src.utils.decorators import handle_errors
from src.utils.log_formatter import (
    OperationLogger,
    log_index_operation,
    LogFormatter
)


class BM25StoreManager:
    """
    BM25存储管理器（单例模式）
    
    确保全局只有一个BM25检索器实例，解决多个组件各自创建独立实例导致状态不共享的问题。
    """
    
    _instance: Optional["BM25StoreManager"] = None
    _initialized: bool = False
    
    def __new__(cls, config: Optional[ConfigManager] = None):
        """实现单例模式"""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(self, config: Optional[ConfigManager] = None):
        """
        初始化BM25存储管理器
        
        Args:
            config: 配置实例，如果为None则使用全局配置
        """
        # 避免重复初始化
        if BM25StoreManager._initialized:
            return
        
        self.config = config or get_config_manager()
        self.bm25_retriever: Optional[ChineseBM25Retriever] = None
        self.docstore: Optional[SimpleDocumentStore] = None
        self.nodes: Optional[List] = None
        self._tokenizer = None
        
        BM25StoreManager._initialized = True
        logger.debug("BM25StoreManager 单例实例已创建")
    
    @classmethod
    def get_instance(cls, config: Optional[ConfigManager] = None) -> "BM25StoreManager":
        """
        获取单例实例
        
        Args:
            config: 配置实例
            
        Returns:
            BM25StoreManager 单例实例
        """
        if cls._instance is None:
            cls._instance = cls(config)
        return cls._instance
    
    @classmethod
    def reset_instance(cls):
        """
        重置单例实例（仅用于测试）
        """
        cls._instance = None
        cls._initialized = False
        logger.debug("BM25StoreManager 单例实例已重置")
    
    def _get_tokenizer(self):
        """获取或创建分词器"""
        if self._tokenizer is None:
            retrieval_config = self.config.retrieval
            self._tokenizer = create_tokenizer(
                language=retrieval_config.bm25_language,
                use_stemmer=(retrieval_config.bm25_language == "english")
            )
        return self._tokenizer
    
    @handle_errors(
        operation_name="初始化/加载 BM25检索器",
        return_on_error=None,
        log_traceback=True,
        suppress_exception=True
    )
    def setup_bm25(self, nodes: Optional[List] = None) -> Optional[ChineseBM25Retriever]:
        """
        初始化或重新加载BM25检索器
        
        Args:
            nodes: 文档节点列表（可选，如果提供则用于创建新的BM25检索器）
        
        Returns:
            ChineseBM25Retriever实例或None
        """
        # 检查BM25Retriever是否可用
        if not HAS_BM25_RETRIEVER or ChineseBM25Retriever is None:
            logger.warning(
                LogFormatter.format_warning(
                    "ChineseBM25Retriever不可用",
                    "跳过BM25检索器初始化"
                )
            )
            return None
        
        # 获取分词器
        tokenizer = self._get_tokenizer()
        bm25_storage_dir = self.config.paths.bm25_storage_dir
        
        if nodes is not None:
            # 创建新的中文BM25检索器
            operation_logger = OperationLogger("初始化中文 BM25 检索器")
            operation_logger.log_start("使用 jieba 分词")
            
            # 保存节点
            self.nodes = nodes
            
            # 创建中文BM25检索器
            self.bm25_retriever = ChineseBM25Retriever(
                nodes=nodes,
                tokenizer=tokenizer,
                similarity_top_k=self.config.retrieval.bm25_top_k,
                verbose=True,
            )
            
            operation_logger.log_success(f"文档数: {len(nodes)}")
        else:
            # 从持久化存储加载
            operation_logger = OperationLogger("加载 BM25 索引")
            operation_logger.log_start(f"从 {bm25_storage_dir}")
            
            # 使用 from_persist_dir 方法加载
            self.bm25_retriever = ChineseBM25Retriever.from_persist_dir(
                persist_dir=str(bm25_storage_dir),
                tokenizer=tokenizer,
                verbose=True
            )
            
            if self.bm25_retriever is not None:
                operation_logger.log_success()
            else:
                operation_logger.log_warning(
                    "BM25 索引加载失败",
                    "将在查询时从向量索引重建"
                )
                return None
        
        return self.bm25_retriever
    
    def persist_bm25(self) -> bool:
        """
        持久化BM25检索器
        
        Returns:
            bool: 是否成功
        """
        if self.bm25_retriever is None:
            logger.warning("BM25检索器未初始化，无法持久化")
            return False
        
        bm25_storage_dir = self.config.paths.bm25_storage_dir
        try:
            logger.info(f"保存 BM25 索引到: {bm25_storage_dir}")
            self.bm25_retriever.persist(str(bm25_storage_dir))
            logger.info("✅ BM25索引已保存")
            return True
        except Exception as e:
            logger.error(f"保存 BM25 索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False
    
    def get_bm25_retriever(self) -> Optional[ChineseBM25Retriever]:
        """
        获取BM25检索器实例
        
        Returns:
            ChineseBM25Retriever实例或None
        """
        return self.bm25_retriever
    
    def is_loaded(self) -> bool:
        """
        检查BM25检索器是否已加载
        
        Returns:
            bool: 是否已加载
        """
        return self.bm25_retriever is not None
    
    def delete_bm25_index(self) -> bool:
        """
        删除BM25索引
        
        Returns:
            bool: 是否成功
        """
        bm25_storage_dir = self.config.paths.bm25_storage_dir
        if not bm25_storage_dir.exists():
            logger.warning("BM25索引目录不存在")
            return True
        
        try:
            import shutil
            logger.info(f"删除 BM25 索引: {bm25_storage_dir}")
            shutil.rmtree(bm25_storage_dir)
            bm25_storage_dir.mkdir(parents=True, exist_ok=True)
            logger.info("✅ BM25 索引已删除")
            self.bm25_retriever = None
            return True
        except Exception as e:
            logger.error(f"删除 BM25 索引失败: {e}")
            return False
