"""
最终修复BM25索引问题
"""
import sys
import os
import json
import shutil
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_params_index_json(bm25_dir: Path):
    """创建缺失的params.index.json文件"""
    params_file = bm25_dir / "params.index.json"
    
    if not params_file.exists():
        print(f"创建 {params_file.name}...")
        
        # 基本的BM25参数
        params = {
            "version": "0.2.14",
            "index_type": "csc",
            "num_docs": 10,  # 估计值
            "num_terms": 100,  # 估计值
            "avg_doc_len": 50,  # 估计值
            "k1": 1.5,
            "b": 0.75,
            "epsilon": 0.25
        }
        
        # 尝试从现有文件推断
        corpus_file = bm25_dir / "corpus.jsonl"
        if corpus_file.exists():
            try:
                with open(corpus_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                    params["num_docs"] = len(lines)
                    
                    # 估计平均文档长度
                    total_tokens = 0
                    for line in lines:
                        try:
                            tokens = json.loads(line.strip())
                            if isinstance(tokens, list):
                                total_tokens += len(tokens)
                        except:
                            pass
                    
                    if len(lines) > 0:
                        params["avg_doc_len"] = total_tokens / len(lines)
                        
                        # 估计词汇表大小
                        unique_tokens = set()
                        for line in lines[:100]:  # 只检查前100行
                            try:
                                tokens = json.loads(line.strip())
                                if isinstance(tokens, list):
                                    unique_tokens.update(tokens)
                            except:
                                pass
                        params["num_terms"] = len(unique_tokens)
            except Exception as e:
                print(f"   警告: 无法从corpus.jsonl推断参数: {e}")
        
        with open(params_file, 'w', encoding='utf-8') as f:
            json.dump(params, f, indent=2)
        
        print(f"   ✅ 已创建")
        return True
    return False

def fix_bm25_index():
    """修复BM25索引"""
    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 False
    
    # 列出文件
    files = list(bm25_dir.glob("*"))
    print(f"   找到 {len(files)} 个文件:")
    for f in files:
        print(f"   - {f.name}")
    
    # 检查缺失的文件
    missing_files = []
    required_files = [
        "params.index.json",
        "corpus.jsonl",
        "corpus.mmindex.json",
        "data.csc.index.npy",
        "indices.csc.index.npy",
        "indptr.csc.index.npy"
    ]
    
    for req_file in required_files:
        if not (bm25_dir / req_file).exists():
            missing_files.append(req_file)
    
    if missing_files:
        print(f"\n2. 缺失的文件: {missing_files}")
        
        # 尝试创建缺失的params.index.json
        if "params.index.json" in missing_files:
            create_params_index_json(bm25_dir)
            missing_files.remove("params.index.json")
        
        # 如果还有其他缺失文件，建议删除并重建
        if missing_files:
            print(f"   仍然缺失文件: {missing_files}")
            print(f"   建议删除并重建BM25索引")
            
            response = input(f"\n   是否删除现有BM25索引并重建? (y/n): ")
            if response.lower() == 'y':
                print(f"   删除BM25索引...")
                try:
                    shutil.rmtree(bm25_dir)
                    bm25_dir.mkdir(parents=True, exist_ok=True)
                    print(f"   ✅ 已删除")
                    
                    # 现在需要重建索引，但这需要文档
                    print(f"\n   注意: BM25索引已删除")
                    print(f"   系统将在下次查询时自动重建BM25索引")
                    print(f"   或者可以运行: python cli.py build-index --data-dir ./my_knowledge_base")
                    
                    return True
                except Exception as e:
                    print(f"   ❌ 删除失败: {e}")
                    return False
            else:
                print(f"   跳过删除")
                return False
    else:
        print(f"\n2. 所有必需文件都存在")
        
        # 尝试创建nodes.json和retriever_config.json
        print(f"\n3. 创建ChineseBM25Retriever所需的文件...")
        
        nodes_file = bm25_dir / "nodes.json"
        config_file = bm25_dir / "retriever_config.json"
        
        if not nodes_file.exists():
            print(f"   创建 nodes.json...")
            
            # 从corpus.jsonl读取
            corpus_file = bm25_dir / "corpus.jsonl"
            nodes_data = []
            
            if corpus_file.exists():
                with open(corpus_file, 'r', encoding='utf-8') as f:
                    for i, line in enumerate(f):
                        try:
                            tokens = json.loads(line.strip())
                            text = " ".join(tokens) if isinstance(tokens, list) else str(tokens)
                            
                            node_dict = {
                                "id_": f"node_{i}",
                                "text": text,
                                "metadata": {
                                    "source": "bm25_corpus",
                                    "index": i,
                                    "tokens_count": len(tokens) if isinstance(tokens, list) else 1
                                }
                            }
                            nodes_data.append(node_dict)
                        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不存在，创建空文件")
                with open(nodes_file, 'w', encoding='utf-8') as f:
                    json.dump([], f, ensure_ascii=False, indent=2)
        
        if not config_file.exists():
            print(f"   创建 retriever_config.json...")
            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"   ✅ 已创建")
        
        print(f"\n4. 修复完成!")
        print(f"   现在ChineseBM25Retriever应该可以正常加载了")
        return True
    
    return False

def test_fix():
    """测试修复是否有效"""
    print(f"\n5. 测试修复...")
    
    try:
        from src.storage import BM25StoreManager
        from src.tokenizers import create_tokenizer
        
        config = get_config_manager()
        bm25_manager = BM25StoreManager(config=config)
        
        # 尝试加载
        tokenizer = create_tokenizer(
            language=config.retrieval.bm25_language,
            use_stemmer=(config.retrieval.bm25_language == "english")
        )
        
        from src.retrievers import ChineseBM25Retriever
        bm25_retriever = ChineseBM25Retriever.from_persist_dir(
            persist_dir=str(config.paths.bm25_storage_dir),
            tokenizer=tokenizer,
            verbose=True
        )
        
        if bm25_retriever:
            print(f"   ✅ BM25检索器加载成功!")
            
            # 简单测试
            test_query = "测试"
            nodes = bm25_retriever.retrieve(test_query)
            print(f"   测试查询 '{test_query}' 检索到 {len(nodes)} 个结果")
            
            return True
        else:
            print(f"   ❌ BM25检索器加载失败")
            return False
            
    except Exception as e:
        print(f"   ❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    if fix_bm25_index():
        test_fix()
    
    print(f"\n" + "="*50)
    print(f"修复完成!")