"""
支持中文的 BM25 检索器
使用 jieba 分词替代默认的 bm25s.tokenize
"""
from typing import List, Optional, Callable
from pathlib import Path
import json
from llama_index.core.schema import NodeWithScore, QueryBundle, BaseNode, TextNode
from llama_index.core.retrievers import BaseRetriever
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
from loguru import logger
import bm25s
import numpy as np


class ChineseBM25Retriever(BaseRetriever):
    """支持中文的 BM25 检索器"""
    
    def __init__(
        self,
        nodes: List[BaseNode],
        tokenizer: Callable[[str], List[str]],
        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
        callback_manager: Optional[CallbackManager] = None,
        verbose: bool = False,
    ):
        """
        初始化中文 BM25 检索器
        
        Args:
            nodes: 文档节点列表
            tokenizer: 分词函数（如 jieba tokenizer）
            similarity_top_k: 返回的top-k结果数量
            callback_manager: 回调管理器
            verbose: 是否显示详细信息
        """
        self._nodes = nodes
        self._tokenizer = tokenizer
        self.similarity_top_k = similarity_top_k
        self._verbose = verbose
        
        # 构建语料库和索引
        logger.info(f"使用自定义 jieba 分词构建 BM25 索引...")
        self._build_index()
        logger.info(f"✅ BM25 索引构建完成 - 文档数: {len(nodes)}")
        
        super().__init__(callback_manager=callback_manager)
    
    def _build_index(self):
        """使用 jieba 分词构建 BM25 索引"""
        # 对所有文档进行分词
        corpus_texts = [node.get_content() for node in self._nodes]
        
        if self._verbose:
            logger.info(f"对 {len(corpus_texts)} 个文档进行分词...")
        
        # 使用 jieba 分词
        corpus_tokens = [self._tokenizer(text) for text in corpus_texts]
        
        if self._verbose:
            # 显示前几个文档的分词示例
            for i, (text, tokens) in enumerate(zip(corpus_texts[:2], corpus_tokens[:2])):
                logger.debug(f"文档 {i}: {text[:50]}...")
                logger.debug(f"  分词: {tokens[:10]}...")
        
        # 创建 BM25 索引
        self.bm25 = bm25s.BM25()
        
        # 将分词结果转换为 bm25s 需要的格式
        # bm25s 需要 Tokenized 对象，但我们可以直接用分词后的文本列表
        self.bm25.index(corpus_tokens, show_progress=self._verbose)
        
        # 保存原始文本和tokens供调试
        self._corpus_texts = corpus_texts
        self._corpus_tokens = corpus_tokens
    
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """
        检索相关文档
        
        Args:
            query_bundle: 查询包
            
        Returns:
            带分数的节点列表
        """
        query = query_bundle.query_str
        
        # 使用 jieba 对查询进行分词
        query_tokens = self._tokenizer(query)
        
        if self._verbose:
            logger.debug(f"查询: {query}")
            logger.debug(f"查询分词: {query_tokens}")
        
        # 执行 BM25 检索
        # bm25s.retrieve 需要列表的列表格式 (list of lists)
        try:
            # 计算k值，确保不超过文档数量
            k_value = min(self.similarity_top_k, len(self._nodes))
            
            # 如果k_value为0，直接返回空结果
            if k_value <= 0:
                return []
            
            # 直接传递列表的列表格式
            results, scores = self.bm25.retrieve(
                [query_tokens],  # 必须是列表的列表
                k=k_value,
                show_progress=False,
            )
        except TypeError as e:
            # 处理TypeError，可能是num_docs为None的问题
            if "'>' not supported between instances of 'int' and 'NoneType'" in str(e):
                logger.warning(f"BM25索引可能未正确初始化，尝试使用文档数量作为k值")
                # 尝试使用文档数量
                try:
                    results, scores = self.bm25.retrieve(
                        [query_tokens],
                        k=len(self._nodes),
                        show_progress=False,
                    )
                except Exception as e2:
                    logger.error(f"BM25 检索失败（第二次尝试）: {e2}")
                    return []
            else:
                logger.error(f"BM25 检索失败: {e}")
                import traceback
                logger.debug(traceback.format_exc())
                return []
        except Exception as e:
            logger.error(f"BM25 检索失败: {e}")
            import traceback
            logger.debug(traceback.format_exc())
            return []
        
        # 构建结果
        nodes_with_scores = []
        
        # results 和 scores 是批量结果，取第一个
        if len(results) > 0 and len(scores) > 0:
            indices = results[0]
            score_values = scores[0]
            
            if self._verbose:
                logger.debug(f"检索到 {len(indices)} 个结果")
                logger.debug(f"分数: {score_values}")
            
            for idx, score in zip(indices, score_values):
                if idx < len(self._nodes):
                    node = self._nodes[idx]
                    nodes_with_scores.append(
                        NodeWithScore(node=node, score=float(score))
                    )
        
        return nodes_with_scores
    
    def persist(self, persist_dir: str):
        """
        持久化索引
        
        Args:
            persist_dir: 持久化目录
        """
        persist_path = Path(persist_dir)
        persist_path.mkdir(parents=True, exist_ok=True)
        
        try:
            # 1. 保存 BM25 索引 (使用 bm25s 内置方法)
            self.bm25.save(persist_dir, corpus=self._corpus_tokens)
            logger.info(f"BM25 索引已保存到: {persist_dir}")
            
            # 2. 保存节点信息 (用于重建时恢复节点)
            nodes_data = []
            for node in self._nodes:
                node_dict = {
                    "id_": node.id_,
                    "text": node.get_content(),
                    "metadata": node.metadata if hasattr(node, 'metadata') else {},
                }
                # 保存额外的节点属性
                if hasattr(node, 'embedding') and node.embedding:
                    node_dict["embedding"] = node.embedding
                nodes_data.append(node_dict)
            
            nodes_file = persist_path / "nodes.json"
            with open(nodes_file, "w", encoding="utf-8") as f:
                json.dump(nodes_data, f, ensure_ascii=False, indent=2)
            logger.info(f"节点信息已保存: {nodes_file}")
            
            # 3. 保存检索器配置
            config_data = {
                "similarity_top_k": self.similarity_top_k,
                "verbose": self._verbose,
                "num_nodes": len(self._nodes),
            }
            config_file = persist_path / "retriever_config.json"
            with open(config_file, "w", encoding="utf-8") as f:
                json.dump(config_data, f, ensure_ascii=False, indent=2)
            logger.info(f"检索器配置已保存: {config_file}")
            
            logger.info("✅ ChineseBM25Retriever 持久化完成")
            
        except Exception as e:
            logger.error(f"持久化 BM25 索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            raise
    
    @classmethod
    def from_persist_dir(
        cls,
        persist_dir: str,
        tokenizer: Callable[[str], List[str]],
        verbose: bool = False,
    ) -> Optional["ChineseBM25Retriever"]:
        """
        从持久化目录加载
        
        Args:
            persist_dir: 持久化目录
            tokenizer: 分词函数（必需，用于查询时分词）
            verbose: 是否显示详细信息
            
        Returns:
            ChineseBM25Retriever 实例，加载失败时返回 None
        """
        persist_path = Path(persist_dir)
        
        # 检查必要的文件是否存在
        nodes_file = persist_path / "nodes.json"
        config_file = persist_path / "retriever_config.json"
        
        # 如果缺少必要文件，尝试从其他文件重建
        if not nodes_file.exists() or not config_file.exists():
            logger.warning(f"缺少必要文件，尝试从旧格式BM25索引重建...")
            return cls._load_from_legacy_format(persist_dir, tokenizer, verbose)
        
        try:
            # 1. 加载节点信息
            with open(nodes_file, "r", encoding="utf-8") as f:
                nodes_data = json.load(f)
            
            # 重建节点对象
            nodes = []
            for node_dict in nodes_data:
                node = TextNode(
                    id_=node_dict.get("id_"),
                    text=node_dict.get("text", ""),
                    metadata=node_dict.get("metadata", {}),
                )
                if "embedding" in node_dict and node_dict["embedding"]:
                    node.embedding = node_dict["embedding"]
                nodes.append(node)
            
            logger.info(f"已加载 {len(nodes)} 个节点")
            
            # 2. 加载检索器配置
            with open(config_file, "r", encoding="utf-8") as f:
                config_data = json.load(f)
            
            similarity_top_k = config_data.get("similarity_top_k", DEFAULT_SIMILARITY_TOP_K)
            
            # 3. 创建新的BM25检索器（重新构建索引）
            # 而不是尝试加载可能损坏的索引
            logger.info(f"重新构建BM25索引...")
            instance = cls(
                nodes=nodes,
                tokenizer=tokenizer,
                similarity_top_k=similarity_top_k,
                verbose=verbose
            )
            
            logger.info(f"✅ ChineseBM25Retriever 从 {persist_dir} 重新构建成功")
            return instance
            
        except Exception as e:
            logger.error(f"从 {persist_dir} 加载 BM25 检索器失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return None
    
    @classmethod
    def _load_from_legacy_format(
        cls,
        persist_dir: str,
        tokenizer: Callable[[str], List[str]],
        verbose: bool = False,
    ) -> Optional["ChineseBM25Retriever"]:
        """
        从旧格式BM25索引加载（兼容性支持）
        
        旧格式可能只有 bm25s 索引文件，没有 nodes.json 和 retriever_config.json
        """
        persist_path = Path(persist_dir)
        
        try:
            # 1. 尝试加载 BM25 索引
            bm25 = bm25s.BM25.load(persist_dir, load_corpus=False)
            logger.info(f"BM25 索引加载成功（旧格式）: {persist_dir}")
            
            # 2. 尝试从 corpus.jsonl 读取文档内容
            corpus_file = persist_path / "corpus.jsonl"
            nodes = []
            
            if corpus_file.exists():
                logger.info(f"从 {corpus_file} 读取文档内容...")
                with open(corpus_file, 'r', encoding='utf-8') as f:
                    for i, line in enumerate(f):
                        try:
                            # corpus.jsonl 可能包含分词后的token列表
                            tokens = json.loads(line.strip())
                            # 将tokens重新组合成文本
                            text = " ".join(tokens) if isinstance(tokens, list) else str(tokens)
                            
                            node = TextNode(
                                id_=f"legacy_node_{i}",
                                text=text,
                                metadata={"source": "legacy_bm25_index", "index": i},
                            )
                            nodes.append(node)
                        except json.JSONDecodeError:
                            continue
                
                logger.info(f"从旧格式索引重建了 {len(nodes)} 个节点")
            else:
                # 如果没有corpus.jsonl，创建虚拟节点
                logger.warning(f"corpus.jsonl 不存在，创建虚拟节点")
                # 尝试从其他文件推断节点数量
                # 检查是否有其他索引文件可以推断大小
                import numpy as np
                indices_file = persist_path / "indices.csc.index.npy"
                if indices_file.exists():
                    indices = np.load(indices_file)
                    # 这是一个粗略的估计
                    estimated_nodes = indices.shape[0] if hasattr(indices, 'shape') else 10
                else:
                    estimated_nodes = 10
                
                for i in range(estimated_nodes):
                    node = TextNode(
                        id_=f"virtual_node_{i}",
                        text=f"虚拟文档 {i} - 从旧BM25索引重建",
                        metadata={"source": "virtual_reconstruction", "index": i},
                    )
                    nodes.append(node)
                
                logger.info(f"创建了 {len(nodes)} 个虚拟节点")
            
            if not nodes:
                logger.error("无法从旧格式索引重建节点")
                return None
            
            # 3. 使用默认配置
            similarity_top_k = DEFAULT_SIMILARITY_TOP_K
            
            # 4. 创建实例
            instance = object.__new__(cls)
            instance._nodes = nodes
            instance._tokenizer = tokenizer
            instance.similarity_top_k = similarity_top_k
            instance._verbose = verbose
            instance.bm25 = bm25
            instance._corpus_texts = [node.get_content() for node in nodes]
            
            # 重新分词生成语料库 tokens
            instance._corpus_tokens = [tokenizer(text) for text in instance._corpus_texts]
            
            # 初始化父类
            BaseRetriever.__init__(instance, callback_manager=None)
            
            logger.info(f"✅ ChineseBM25Retriever 从旧格式索引加载成功")
            logger.info(f"   节点数: {len(nodes)}")
            logger.info(f"   注意: 这是兼容性加载，建议重新构建索引以获得最佳效果")
            
            # 保存新的格式以便下次使用
            try:
                instance.persist(persist_dir)
                logger.info(f"已将旧格式索引转换为新格式")
            except Exception as e:
                logger.warning(f"无法保存新格式索引: {e}")
            
            return instance
            
        except Exception as e:
            logger.error(f"从旧格式索引加载失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return None
