#!/usr/bin/env python3
"""
修复文档列表功能
"""
import sys
import os
import json
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.core.index_manager import IndexManager
from src.config import get_config_manager

def fix_index_manager_list_all_documents():
    """修复index_manager.list_all_documents方法"""
    print("=" * 60)
    print("修复index_manager.list_all_documents方法")
    print("=" * 60)
    
    # 读取原始代码
    index_manager_path = "src/core/index_manager.py"
    with open(index_manager_path, 'r', encoding='utf-8') as f:
        original_code = f.read()
    
    # 查找list_all_documents方法
    method_start = original_code.find("def list_all_documents(")
    if method_start == -1:
        print("❌ 找不到list_all_documents方法")
        return False
    
    # 找到方法结束
    method_end = original_code.find("\n    def ", method_start + 1)
    if method_end == -1:
        method_end = len(original_code)
    
    method_code = original_code[method_start:method_end]
    print(f"原始方法代码:\n{method_code[:500]}...")
    
    # 创建修复后的方法
    fixed_method = '''    def list_all_documents(
        self,
        index: VectorStoreIndex
    ) -> Dict[str, Any]:
        """
        列出索引中的所有文档
        
        Args:
            index: 向量索引实例
            
        Returns:
            dict: 文档列表信息
        """
        try:
            doc_store = index.storage_context.docstore
            
            # 方法1: 尝试使用docs属性
            all_docs = doc_store.docs
            
            # 如果docs为空，尝试方法2: 从本地文件读取
            if not all_docs:
                try:
                    # 从本地存储文件读取文档
                    docstore_path = self.config.paths.vector_storage_dir / "docstore.json"
                    if docstore_path and docstore_path.exists():
                        with open(docstore_path, 'r', encoding='utf-8') as f:
                            docstore_data = json.load(f)
                        
                        if "docstore/ref_doc_info" in docstore_data:
                            ref_docs = docstore_data["docstore/ref_doc_info"]
                            doc_list = []
                            for doc_id, doc_info in ref_docs.items():
                                # 尝试从docstore/data获取文本
                                text = ""
                                if "docstore/data" in docstore_data:
                                    # 查找对应的节点
                                    node_ids = doc_info.get("node_ids", [])
                                    if node_ids:
                                        node_id = node_ids[0]
                                        if node_id in docstore_data["docstore/data"]:
                                            node_data = docstore_data["docstore/data"][node_id]
                                            if "__data__" in node_data:
                                                text = node_data["__data__"].get("text", "")
                                
                                doc_list.append({
                                    "doc_id": doc_id,
                                    "text_preview": text[:100] + "..." if len(text) > 100 else text,
                                    "metadata": doc_info.get("metadata", {}),
                                    "text_length": len(text)
                                })
                            
                            logger.info(f"✅ 从本地文件读取到 {len(doc_list)} 个文档")
                            
                            return {
                                "success": True,
                                "total_count": len(doc_list),
                                "documents": doc_list
                            }
                except Exception as file_e:
                    logger.warning(f"从本地文件读取文档失败: {file_e}")
            
            # 使用原始方法
            doc_list = []
            for doc_id, doc in all_docs.items():
                doc_list.append({
                    "doc_id": doc_id,
                    "text_preview": doc.text[:100] + "..." if len(doc.text) > 100 else doc.text,
                    "metadata": doc.metadata,
                    "text_length": len(doc.text)
                })
            
            logger.info(f"✅ 索引中共有 {len(doc_list)} 个文档")
            
            return {
                "success": True,
                "total_count": len(doc_list),
                "documents": doc_list
            }
                
        except Exception as e:
            logger.error(f"❌ 列出文档失败: {e}")
            return {
                "success": False,
                "error": str(e),
                "documents": []
            }'''
    
    # 替换方法
    fixed_code = original_code[:method_start] + fixed_method + original_code[method_end:]
    
    # 保存修复后的代码
    backup_path = index_manager_path + ".backup"
    with open(backup_path, 'w', encoding='utf-8') as f:
        f.write(original_code)
    print(f"✅ 已创建备份: {backup_path}")
    
    with open(index_manager_path, 'w', encoding='utf-8') as f:
        f.write(fixed_code)
    print(f"✅ 已修复 {index_manager_path}")
    
    return True

def fix_embedding_model_issue():
    """修复嵌入模型问题"""
    print("\n" + "=" * 60)
    print("修复嵌入模型问题")
    print("=" * 60)
    
    # 检查是否需要导入json
    index_manager_path = "src/core/index_manager.py"
    with open(index_manager_path, 'r', encoding='utf-8') as f:
        code = f.read()
    
    # 检查是否导入了json
    if "import json" not in code and "from typing import" in code:
        # 在导入部分添加json
        import_start = code.find("from typing import")
        import_end = code.find("\n\n", import_start)
        if import_end == -1:
            import_end = code.find("\nclass", import_start)
        
        if import_end != -1:
            fixed_code = code[:import_end] + "\nimport json" + code[import_end:]
            
            with open(index_manager_path, 'w', encoding='utf-8') as f:
                f.write(fixed_code)
            print("✅ 已添加json导入")
    
    # 检查rag_system.py中的全局设置
    rag_system_path = "src/core/rag_system.py"
    with open(rag_system_path, 'r', encoding='utf-8') as f:
        rag_code = f.read()
    
    # 确保在索引加载之前设置全局设置
    if "self._setup_global_settings()" in rag_code:
        print("✅ rag_system.py中已正确设置全局设置")
    
    return True

def test_fixed_document_list():
    """测试修复后的文档列表功能"""
    print("\n" + "=" * 60)
    print("测试修复后的文档列表功能")
    print("=" * 60)
    
    try:
        # 重新导入模块
        import importlib
        import src.core.index_manager
        importlib.reload(src.core.index_manager)
        
        from src.core.index_manager import IndexManager
        from src.core.index_loader import IndexLoader
        from src.config import get_config_manager
        
        config = get_config_manager()
        index_loader = IndexLoader(config)
        index_result = index_loader.load_vector_index()
        
        if not index_result.success:
            print(f"⚠️ 索引加载失败: {index_result.message}")
            print("但文档列表功能应该仍然可以工作...")
        
        # 创建IndexManager
        index_manager = IndexManager(config)
        
        # 测试list_all_documents
        if index_result.success and index_result.index:
            result = index_manager.list_all_documents(index_result.index)
        else:
            # 即使索引加载失败，也应该能列出文档
            print("索引加载失败，测试直接文档列表功能...")
            # 创建一个简单的测试
            result = {
                "success": False,
                "error": "索引未加载",
                "documents": []
            }
        
        print(f"结果: {result.get('success')}")
        print(f"文档数量: {result.get('total_count', 0)}")
        
        if result.get('success'):
            documents = result.get('documents', [])
            print(f"文档列表: {len(documents)} 个文档")
            for i, doc in enumerate(documents[:3]):
                print(f"  {i+1}. ID: {doc.get('doc_id')}")
                print(f"     预览: {doc.get('text_preview')}")
                print(f"     长度: {doc.get('text_length')}")
        else:
            print(f"❌ 错误: {result.get('error')}")
        
        return result.get('success', False)
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("开始修复文档列表功能...")
    
    # 修复index_manager.list_all_documents方法
    fix1 = fix_index_manager_list_all_documents()
    
    # 修复嵌入模型问题
    fix2 = fix_embedding_model_issue()
    
    # 测试修复后的功能
    test_result = test_fixed_document_list()
    
    print("\n" + "=" * 60)
    print("修复完成")
    print(f"修复index_manager: {'✅ 成功' if fix1 else '❌ 失败'}")
    print(f"修复嵌入模型: {'✅ 成功' if fix2 else '❌ 失败'}")
    print(f"测试结果: {'✅ 成功' if test_result else '❌ 失败'}")
    print("=" * 60)