"""
简单重建BM25索引
直接使用现有的BM25索引文件，添加缺失的配置文件
"""
import sys
import os
import json
from pathlib import Path

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.tokenizers import create_tokenizer

def create_missing_config_files(bm25_dir: Path):
    """创建缺失的配置文件"""
    
    # 1. 创建nodes.json文件
    nodes_file = bm25_dir / "nodes.json"
    if not nodes_file.exists():
        print(f"创建 {nodes_file.name}...")
        
        # 从corpus.jsonl读取文档
        corpus_file = bm25_dir / "corpus.jsonl"
        if corpus_file.exists():
            nodes_data = []
            with open(corpus_file, 'r', encoding='utf-8') as f:
                for i, line in enumerate(f):
                    try:
                        doc = json.loads(line.strip())
                        node_data = {
                            "id_": f"node_{i}",
                            "text": doc.get("text", ""),
                            "metadata": doc.get("metadata", {})
                        }
                        nodes_data.append(node_data)
                    except json.JSONDecodeError:
                        continue
            
            with open(nodes_file, 'w', encoding='utf-8') as f:
                json.dump(nodes_data, f, ensure_ascii=False, indent=2)
            print(f"   ✅ 已创建，包含 {len(nodes_data)} 个节点")
        else:
            print(f"   ⚠️ corpus.jsonl 不存在，创建空的 nodes.json")
            with open(nodes_file, 'w', encoding='utf-8') as f:
                json.dump([], f, ensure_ascii=False, indent=2)
    
    # 2. 创建retriever_config.json文件
    config_file = bm25_dir / "retriever_config.json"
    if not config_file.exists():
        print(f"创建 {config_file.name}...")
        config_data = {
            "similarity_top_k": 5,
            "verbose": True,
            "num_nodes": 0
        }
        
        # 尝试从nodes.json获取节点数量
        if nodes_file.exists():
            with open(nodes_file, 'r', encoding='utf-8') as f:
                nodes_data = json.load(f)
                config_data["num_nodes"] = len(nodes_data)
        
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config_data, f, ensure_ascii=False, indent=2)
        print(f"   ✅ 已创建")
    
    return nodes_file.exists() and config_file.exists()

def main():
    print("=== 简单修复 BM25 索引 ===\n")
    
    # 获取配置
    config = get_config_manager()
    
    # 创建BM25存储管理器
    bm25_manager = BM25StoreManager(config=config)
    
    # 检查BM25存储目录
    bm25_dir = config.paths.bm25_storage_dir
    print(f"1. BM25存储目录: {bm25_dir}")
    
    if not bm25_dir.exists():
        print("   ❌ 目录不存在")
        return
    
    # 列出文件
    files = list(bm25_dir.glob("*"))
    print(f"   找到 {len(files)} 个文件:")
    for f in files:
        print(f"   - {f.name}")
    
    # 检查并创建缺失的配置文件
    print("\n2. 检查并创建缺失的配置文件...")
    if create_missing_config_files(bm25_dir):
        print("   ✅ 配置文件已创建/更新")
    else:
        print("   ⚠️ 配置文件创建失败")
    
    # 尝试加载BM25检索器
    print("\n3. 尝试加载BM25检索器...")
    
    # 创建分词器
    tokenizer = create_tokenizer(
        language=config.retrieval.bm25_language,
        use_stemmer=(config.retrieval.bm25_language == "english")
    )
    
    # 尝试从持久化目录加载
    try:
        from src.retrievers import ChineseBM25Retriever
        bm25_retriever = ChineseBM25Retriever.from_persist_dir(
            persist_dir=str(bm25_dir),
            tokenizer=tokenizer,
            verbose=True
        )
        
        if bm25_retriever:
            print("   ✅ BM25检索器加载成功")
            
            # 测试检索
            print("\n4. 测试BM25检索:")
            test_queries = ["什么是RAG", "向量数据库", "人工智能"]
            
            for query in test_queries:
                print(f"\n   查询: {query}")
                nodes = bm25_retriever.retrieve(query)
                print(f"   检索到 {len(nodes)} 个结果")
                
                if nodes:
                    for i, node in enumerate(nodes[:2], 1):
                        score = node.score if hasattr(node, 'score') else 'N/A'
                        text_preview = node.node.get_content()[:80] if hasattr(node.node, 'get_content') else 'N/A'
                        print(f"      [{i}] 分数: {score}")
                        print(f"          内容: {text_preview}...")
                else:
                    print("      ⚠️ 未检索到结果")
        else:
            print("   ❌ BM25检索器加载失败")
            
    except Exception as e:
        print(f"   ❌ 加载失败: {e}")
        import traceback
        traceback.print_exc()
        
        # 尝试另一种方法：直接使用bm25s库
        print("\n5. 尝试直接使用bm25s库...")
        try:
            import bm25s
            bm25 = bm25s.BM25.load(str(bm25_dir), load_corpus=False)
            print("   ✅ bm25s索引加载成功")
            
            # 测试检索
            print("   测试检索...")
            test_query = "什么是RAG"
            
            # 需要分词器
            query_tokens = tokenizer(test_query)
            print(f"   查询分词: {query_tokens}")
            
            results, scores = bm25.retrieve([query_tokens], k=5)
            print(f"   检索结果: {results}")
            print(f"   分数: {scores}")
            
        except Exception as e2:
            print(f"   ❌ bm25s加载失败: {e2}")
    
    print("\n" + "="*50)
    print("修复完成！")

if __name__ == "__main__":
    main()