"""
完整修复BM25索引问题
创建所有缺失的bm25s库所需文件
"""
import sys
import os
import json
import numpy as np
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

def create_bm25s_files(bm25_dir: Path):
    """创建所有bm25s库所需的文件"""
    
    print(f"创建bm25s库所需文件...")
    
    # 1. 检查corpus.jsonl文件
    corpus_file = bm25_dir / "corpus.jsonl"
    if not corpus_file.exists():
        print(f"   ❌ corpus.jsonl不存在，无法修复")
        return False
    
    # 读取corpus.jsonl
    documents = []
    with open(corpus_file, 'r', encoding='utf-8') as f:
        for line in f:
            try:
                tokens = json.loads(line.strip())
                if isinstance(tokens, list):
                    documents.append(tokens)
            except json.JSONDecodeError:
                continue
    
    if not documents:
        print(f"   ⚠️ corpus.jsonl为空或格式错误")
        return False
    
    num_docs = len(documents)
    print(f"   找到 {num_docs} 个文档")
    
    # 2. 创建vocab.index.json
    vocab_file = bm25_dir / "vocab.index.json"
    if not vocab_file.exists():
        print(f"   创建 vocab.index.json...")
        
        # 收集所有唯一的token
        vocab_set = set()
        for doc_tokens in documents:
            vocab_set.update(doc_tokens)
        
        vocab_list = sorted(list(vocab_set))
        vocab_dict = {token: idx for idx, token in enumerate(vocab_list)}
        
        with open(vocab_file, 'w', encoding='utf-8') as f:
            json.dump(vocab_dict, f, ensure_ascii=False, indent=2)
        
        print(f"      ✅ 已创建，词汇表大小: {len(vocab_list)}")
    
    # 3. 创建params.index.json（如果不存在）
    params_file = bm25_dir / "params.index.json"
    if not params_file.exists():
        print(f"   创建 params.index.json...")
        
        # 计算平均文档长度
        total_tokens = sum(len(doc) for doc in documents)
        avg_doc_len = total_tokens / num_docs if num_docs > 0 else 50
        
        params = {
            "version": "0.2.14",
            "index_type": "csc",
            "num_docs": num_docs,
            "num_terms": len(vocab_set) if 'vocab_set' in locals() else 100,
            "avg_doc_len": avg_doc_len,
            "k1": 1.5,
            "b": 0.75,
            "epsilon": 0.25
        }
        
        with open(params_file, 'w', encoding='utf-8') as f:
            json.dump(params, f, indent=2)
        
        print(f"      ✅ 已创建")
    
    # 4. 创建nodes.json和retriever_config.json（用于ChineseBM25Retriever）
    nodes_file = bm25_dir / "nodes.json"
    if not nodes_file.exists():
        print(f"   创建 nodes.json...")
        
        nodes_data = []
        for i, doc_tokens in enumerate(documents):
            text = " ".join(doc_tokens)
            nodes_data.append({
                "id_": f"node_{i}",
                "text": text,
                "metadata": {
                    "source": "bm25_corpus",
                    "index": i,
                    "tokens_count": len(doc_tokens)
                }
            })
        
        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)} 个节点")
    
    config_file = bm25_dir / "retriever_config.json"
    if not config_file.exists():
        print(f"   创建 retriever_config.json...")
        
        config_data = {
            "similarity_top_k": 5,
            "verbose": True,
            "num_nodes": num_docs
        }
        
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config_data, f, ensure_ascii=False, indent=2)
        
        print(f"      ✅ 已创建")
    
    # 5. 检查其他必要的npy文件是否存在
    required_npy_files = [
        "data.csc.index.npy",
        "indices.csc.index.npy", 
        "indptr.csc.index.npy"
    ]
    
    missing_npy = []
    for npy_file in required_npy_files:
        if not (bm25_dir / npy_file).exists():
            missing_npy.append(npy_file)
    
    if missing_npy:
        print(f"   ⚠️ 缺失npy文件: {missing_npy}")
        print(f"   注意: 这些是二进制索引文件，无法手动创建")
        print(f"   可能需要重新构建BM25索引")
        return False
    
    return True

def test_bm25_loading():
    """测试BM25加载"""
    print(f"\n测试BM25加载...")
    
    try:
        from src.tokenizers import create_tokenizer
        from src.retrievers import ChineseBM25Retriever
        
        config = get_config_manager()
        
        # 创建分词器
        tokenizer = create_tokenizer(
            language=config.retrieval.bm25_language,
            use_stemmer=(config.retrieval.bm25_language == "english")
        )
        
        # 尝试加载
        bm25_retriever = ChineseBM25Retriever.from_persist_dir(
            persist_dir=str(config.paths.bm25_storage_dir),
            tokenizer=tokenizer,
            verbose=True
        )
        
        if bm25_retriever:
            print(f"   ✅ ChineseBM25Retriever加载成功!")
            
            # 简单测试
            test_query = "人工智能"
            print(f"   测试查询: '{test_query}'")
            
            nodes = bm25_retriever.retrieve(test_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 = node.node.get_content()[:80] if hasattr(node.node, 'get_content') else str(node.node)[:80]
                    print(f"      [{i}] 分数: {score}")
                    print(f"          内容: {text}...")
            
            return True
        else:
            print(f"   ❌ ChineseBM25Retriever加载失败")
            return False
            
    except Exception as e:
        print(f"   ❌ 加载测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    print("=== 完整修复 BM25 索引 ===\n")
    
    # 获取配置
    config = get_config_manager()
    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}")
    
    # 创建缺失的文件
    if create_bm25s_files(bm25_dir):
        print(f"\n2. 文件创建完成")
        
        # 测试加载
        if test_bm25_loading():
            print(f"\n✅ BM25索引修复成功!")
        else:
            print(f"\n⚠️ BM25索引修复部分成功，但加载测试失败")
            print(f"   可能需要重新构建索引")
    else:
        print(f"\n❌ BM25索引修复失败")
        print(f"   建议删除现有索引并重新构建")
        print(f"   删除命令: Remove-Item -Path '{bm25_dir}' -Recurse -Force")
        print(f"   重建命令: python cli.py build-index --data-dir ./my_knowledge_base")
    
    print(f"\n" + "="*50)
    print(f"修复完成!")

if __name__ == "__main__":
    main()