"""
修复BM25索引加载问题
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from loguru import logger
from src.config import get_config_manager
from src.storage import BM25StoreManager
from src.storage.strategies import StorageProvider

def main():
    print("=== 修复 BM25 索引加载问题 ===\n")
    
    # 获取配置
    config = get_config_manager()
    
    # 创建BM25存储管理器
    bm25_manager = BM25StoreManager(config=config)
    
    # 检查BM25索引是否存在
    bm25_storage_dir = config.paths.bm25_storage_dir
    print(f"1. 检查BM25存储目录: {bm25_storage_dir}")
    
    if bm25_storage_dir.exists():
        files = list(bm25_storage_dir.glob("*"))
        print(f"   找到 {len(files)} 个文件:")
        for f in files:
            print(f"   - {f.name}")
    else:
        print("   ❌ BM25存储目录不存在")
        return
    
    # 检查必要的文件
    nodes_file = bm25_storage_dir / "nodes.json"
    config_file = bm25_storage_dir / "retriever_config.json"
    
    print(f"\n2. 检查必要文件:")
    print(f"   nodes.json 存在: {nodes_file.exists()}")
    print(f"   retriever_config.json 存在: {config_file.exists()}")
    
    # 如果缺少必要文件，需要重建索引
    if not nodes_file.exists() or not config_file.exists():
        print("\n3. 缺少必要文件，需要重建BM25索引")
        
        # 首先需要获取文档节点
        print("   尝试从向量索引获取文档节点...")
        
        # 创建存储提供者
        storage_provider = StorageProvider.get_instance()
        storage_provider.setup(config)
        
        # 获取当前存储策略
        strategy = storage_provider.get_current_strategy()
        
        if strategy is None:
            print("   ❌ 存储策略未初始化")
            return
        
        if strategy.exists():
            print(f"   向量索引存在，后端: {strategy.backend_name}")
            
            # 加载向量索引
            try:
                index = strategy.load_index()
                if index and hasattr(index, 'docstore'):
                    docstore = index.docstore
                    all_nodes = list(docstore.docs.values())
                    print(f"   从向量索引获取到 {len(all_nodes)} 个节点")
                    
                    # 删除旧的BM25索引
                    print("   删除旧的BM25索引...")
                    bm25_manager.delete_bm25_index()
                    
                    # 重建BM25索引
                    print("   使用正确的分词器重建BM25索引...")
                    bm25_retriever = bm25_manager.setup_bm25(nodes=all_nodes)
                    
                    if bm25_retriever:
                        print("   ✅ BM25索引重建成功")
                        
                        # 持久化
                        print("   保存BM25索引...")
                        bm25_manager.persist_bm25()
                        print("   ✅ BM25索引已保存")
                        
                        # 测试检索
                        print("\n4. 测试BM25检索:")
                        test_query = "什么是RAG"
                        print(f"   查询: {test_query}")
                        
                        nodes = bm25_retriever.retrieve(test_query)
                        print(f"   检索到 {len(nodes)} 个结果")
                        
                        if nodes:
                            for i, node in enumerate(nodes[:3], 1):
                                score = node.score if hasattr(node, 'score') else 'N/A'
                                text_preview = node.node.text[:80] if hasattr(node.node, 'text') else 'N/A'
                                print(f"\n   [{i}] 分数: {score}")
                                print(f"       内容: {text_preview}...")
                        else:
                            print("   ⚠️ 未检索到结果")
                    else:
                        print("   ❌ BM25索引重建失败")
                else:
                    print("   ❌ 无法从向量索引获取节点")
            except Exception as e:
                print(f"   ❌ 加载向量索引失败: {e}")
                import traceback
                traceback.print_exc()
        else:
            print("   ❌ 向量索引不存在，请先构建索引")
    else:
        print("\n3. 必要文件存在，尝试加载BM25检索器...")
        bm25_retriever = bm25_manager.setup_bm25(nodes=None)
        
        if bm25_retriever:
            print("   ✅ BM25检索器加载成功")
            
            # 测试检索
            print("\n4. 测试BM25检索:")
            test_query = "什么是RAG"
            print(f"   查询: {test_query}")
            
            nodes = bm25_retriever.retrieve(test_query)
            print(f"   检索到 {len(nodes)} 个结果")
            
            if nodes:
                for i, node in enumerate(nodes[:3], 1):
                    score = node.score if hasattr(node, 'score') else 'N/A'
                    text_preview = node.node.text[:80] if hasattr(node.node, 'text') else 'N/A'
                    print(f"\n   [{i}] 分数: {score}")
                    print(f"       内容: {text_preview}...")
            else:
                print("   ⚠️ 未检索到结果")
        else:
            print("   ❌ BM25检索器加载失败")
    
    print("\n" + "="*50)
    print("修复完成！")

if __name__ == "__main__":
    main()