"""
RAG 系统核心实现（增强版）
基于 LlamaIndex 框架构建索引和检索功能
支持向量检索、BM25检索和混合检索
支持后处理器、元数据过滤、批量查询和增量索引管理
"""
from typing import Optional, List, Dict, Any
from loguru import logger

from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI

# 尝试导入 DeepSeek LLM
try:
    from llama_index.llms.deepseek import DeepSeek as DeepSeekLLM
    HAS_DEEPSEEK_LLM = True
except ImportError:
    HAS_DEEPSEEK_LLM = False
    DeepSeekLLM = None

from src.config import get_config_manager, ConfigManager
from src.embeddings import SiliconFlowEmbedding
from src.utils import setup_logger, apply_all_patches
from src.storage.strategies import StorageProvider
from src.events import FileWatcherService, IncrementalIndexer
from .index_builder import IndexBuilder
from .index_loader import IndexLoader
from .query_engine import QueryEngineManager
from .index_manager import IndexManager


class RAGSystem:
    """RAG 知识库系统核心类"""
    
    def __init__(self, config: Optional[ConfigManager] = None):
        """
        初始化 RAG 系统
        
        Args:
            config: 配置实例，如果为None则使用全局配置
        """
        # 应用所有兼容性补丁（必须在最开始）
        apply_all_patches()
        
        self.config = config or get_config_manager()
        self.index: Optional[VectorStoreIndex] = None
        
        # 初始化存储提供者（单例，所有组件共享）
        self.storage_provider = StorageProvider.get_instance()
        self.storage_provider.setup(self.config)
        
        # 初始化组件
        self.index_builder = IndexBuilder(config=self.config)
        self.index_loader = IndexLoader(config=self.config)
        self.query_engine_manager = QueryEngineManager(config=self.config)
        self.index_manager = IndexManager(config=self.config)
        
        # 配置日志
        setup_logger(self.config.paths.log_file, self.config.log.level)
        
        # 配置全局设置
        self._setup_global_settings()
        
        # 初始化存储策略（如果需要）
        strategy = self.storage_provider.get_current_strategy()
        init_result = strategy.initialize()
        if init_result.success:
            logger.info(f"存储策略已初始化: {strategy.backend_name}")
        
        # 初始化文件监控和增量索引服务
        self._init_file_watcher_and_indexer()
        
        logger.info("RAG 系统初始化完成")
        logger.info(f"检索模式: {self.config.retrieval.mode}")
        logger.info(f"存储后端: {strategy.backend_name}")
    
    def _init_file_watcher_and_indexer(self):
        """初始化文件监控和增量索引服务"""
        try:
            # 检查配置是否启用文件监控
            if not self.config.file_watcher.enabled:
                logger.info("文件监控功能已禁用，跳过初始化")
                return
            
            # 获取知识库目录
            knowledge_dir = str(self.config.paths.knowledge_dir)
            
            # 初始化文件监控服务
            self.file_watcher = FileWatcherService.get_instance()
            
            # 启动文件监控
            watcher_started = self.file_watcher.start(knowledge_dir)
            if watcher_started:
                logger.info(f"✅ 文件监控服务已启动，监控目录: {knowledge_dir}")
            else:
                logger.warning("文件监控服务启动失败")
            
            # 初始化增量索引器
            self.incremental_indexer = IncrementalIndexer.get_instance()
            
            # 启动增量索引器
            indexer_started = self.incremental_indexer.start()
            if indexer_started:
                logger.info("✅ 增量索引器已启动")
            else:
                logger.warning("增量索引器启动失败")
                
            logger.info("文件监控和增量索引服务初始化完成")
            
        except Exception as e:
            logger.error(f"初始化文件监控和增量索引服务失败: {e}")
            logger.warning("文件监控功能将不可用")
    
    def _setup_global_settings(self):
        """配置 LlamaIndex 全局设置"""
        try:
            llm_config = self.config.llm
            embedding_config = self.config.embedding
            retrieval_config = self.config.retrieval
            
            # 设置 LLM
            if HAS_DEEPSEEK_LLM and "deepseek" in llm_config.model.lower():
                logger.info(f"使用 DeepSeek LLM: {llm_config.model}")
                Settings.llm = DeepSeekLLM(
                    model=llm_config.model,
                    temperature=llm_config.temperature,
                    api_key=llm_config.get_api_key(),
                )
            else:
                logger.info(f"使用 OpenAI 兼容 LLM: {llm_config.model}")
                Settings.llm = OpenAI(
                    model=llm_config.model,
                    temperature=llm_config.temperature,
                    api_key=llm_config.get_api_key(),
                    api_base=llm_config.base_url,
                    max_tokens=llm_config.max_tokens,
                    context_window=llm_config.context_window,
                )
            
            # 设置 Embedding 模型
            logger.info(f"使用 Embedding 模型: {embedding_config.model} (via SiliconFlow)")
            Settings.embed_model = SiliconFlowEmbedding(
                api_key=embedding_config.get_api_key(),
                base_url=embedding_config.base_url,
                model_name=embedding_config.model,
            )
            
            # 设置文本分块器
            Settings.node_parser = SentenceSplitter(
                chunk_size=retrieval_config.chunk_size,
                chunk_overlap=retrieval_config.chunk_overlap,
            )
            
            logger.info(f"全局设置配置完成 - LLM: {llm_config.model}, Embedding: {embedding_config.model}")
            
        except Exception as e:
            logger.error(f"配置全局设置失败: {e}")
            raise
    
    def build_index(self, data_dir: str, force_rebuild: bool = False) -> bool:
        """
        构建知识库索引
        
        Args:
            data_dir: 文档目录路径
            force_rebuild: 是否强制重建索引
        
        Returns:
            bool: 构建是否成功
        """
        result = self.index_builder.build_all_indexes(data_dir, force_rebuild)
        if result.success:
            self.index = result.index
        return result.success
    
    def load_index(self) -> bool:
        """
        从存储加载索引
        
        Returns:
            bool: 加载是否成功
        """
        result = self.index_loader.load_all_indexes()
        if result.success:
            self.index = result.index
        return result.success
    
    def query(self, question: str, top_k: Optional[int] = None, mode: Optional[str] = None) -> dict:
        """
        查询知识库
        
        Args:
            question: 用户问题
            top_k: 检索的文档数量
            mode: 检索模式
        
        Returns:
            dict: 包含答案和元数据的字典
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建或加载索引",
                        "answer": None,
                        "sources": [],
                        "mode": None
                    }
            
            # 确定检索模式
            retrieval_mode = mode or self.config.retrieval.mode
            logger.info(f"查询问题: {question}")
            logger.info(f"检索模式: {retrieval_mode}")
            
            # 创建查询引擎
            try:
                query_engine = self.query_engine_manager.create_query_engine(
                    self.index, retrieval_mode, top_k
                )
            except ValueError as e:
                return {
                    "success": False,
                    "error": str(e),
                    "answer": None,
                    "sources": [],
                    "mode": retrieval_mode
                }
            
            # 执行查询
            return self.query_engine_manager.execute_query(query_engine, question, retrieval_mode)
            
        except Exception as e:
            logger.error(f"查询失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e),
                "answer": None,
                "sources": [],
                "mode": mode or self.config.retrieval.mode
            }
    
    def retrieve(
        self,
        question: str,
        top_k: Optional[int] = None,
        mode: Optional[str] = None,
        similarity_cutoff: Optional[float] = None,
        use_reranker: Optional[bool] = None,
        metadata_filters: Optional[Dict] = None
    ) -> dict:
        """
        仅检索相关文档（增强版），支持后处理器和元数据过滤
        
        Args:
            question: 用户问题
            top_k: 检索的文档数量
            mode: 检索模式
            similarity_cutoff: 相似度阈值过滤
            use_reranker: 是否使用重排序
            metadata_filters: 元数据过滤条件
        
        Returns:
            dict: 包含检索结果和元数据的字典
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建或加载索引",
                        "results": [],
                        "mode": None
                    }
            
            # 确定检索模式
            retrieval_mode = mode if mode is not None else "hybrid"
            logger.info(f"检索问题: {question}")
            logger.info(f"检索模式: {retrieval_mode}")
            
            # 创建检索器
            try:
                retriever = self.query_engine_manager.create_retriever(
                    self.index,
                    retrieval_mode,
                    top_k,
                    metadata_filters=metadata_filters
                )
            except ValueError as e:
                return {
                    "success": False,
                    "error": str(e),
                    "results": [],
                    "mode": retrieval_mode
                }
            
            # 执行检索
            return self.query_engine_manager.execute_retrieve(retriever, question, retrieval_mode)
            
        except Exception as e:
            logger.error(f"检索失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e),
                "results": [],
                "mode": mode if mode is not None else "hybrid"
            }
    
    def get_index_stats(self) -> dict:
        """
        获取索引统计信息
        
        Returns:
            dict: 索引统计信息
        """
        return self.index_manager.get_index_stats(self.index)
    
    def delete_index(self) -> bool:
        """
        删除所有索引
        
        Returns:
            bool: 删除是否成功
        """
        success = self.index_manager.delete_index()
        if success:
            self.index = None
        return success
    
    # ============================================================
    # 新增功能：批量查询和增量索引管理
    # ============================================================
    
    def batch_query(
        self,
        questions: List[str],
        mode: Optional[str] = None,
        top_k: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """
        批量异步查询
        
        Args:
            questions: 问题列表
            mode: 检索模式
            top_k: 每个问题检索的文档数量
        
        Returns:
            查询结果列表
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return [{
                        "index": i,
                        "question": q,
                        "success": False,
                        "error": "索引未加载"
                    } for i, q in enumerate(questions, 1)]
            
            # 确定检索模式
            retrieval_mode = mode or self.config.retrieval.mode
            
            # 创建查询引擎
            query_engine = self.query_engine_manager.create_query_engine(
                self.index, retrieval_mode, top_k
            )
            
            # 执行批量查询
            return self.query_engine_manager.batch_query(query_engine, questions)
            
        except Exception as e:
            logger.error(f"批量查询失败: {e}")
            return [{
                "index": i,
                "question": q,
                "success": False,
                "error": str(e)
            } for i, q in enumerate(questions, 1)]
    
    def insert_documents(self, documents: List) -> Dict[str, Any]:
        """
        插入新文档到索引
        
        Args:
            documents: 文档列表
        
        Returns:
            操作结果
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建索引"
                    }
            
            return self.index_manager.insert_documents(self.index, documents)
            
        except Exception as e:
            logger.error(f"插入文档失败: {e}")
            import traceback
            return {
                "success": False,
                "error": str(e),
                "traceback": traceback.format_exc()
            }
    
    def update_documents(self, documents: List) -> Dict[str, Any]:
        """
        更新索引中的文档
        
        Args:
            documents: 文档列表
        
        Returns:
            操作结果
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建索引"
                    }
            
            return self.index_manager.update_documents(self.index, documents)
            
        except Exception as e:
            logger.error(f"更新文档失败: {e}")
            import traceback
            return {
                "success": False,
                "error": str(e),
                "traceback": traceback.format_exc()
            }
    
    def delete_document(self, doc_id: str) -> Dict[str, Any]:
        """
        从索引中删除文档
        
        Args:
            doc_id: 文档ID
        
        Returns:
            操作结果
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建索引"
                    }
            
            return self.index_manager.delete_document(self.index, doc_id)
            
        except Exception as e:
            logger.error(f"删除文档失败: {e}")
            import traceback
            return {
                "success": False,
                "error": str(e),
                "traceback": traceback.format_exc()
            }
    
    def get_document(self, doc_id: str) -> Dict[str, Any]:
        """
        获取索引中的文档
        
        Args:
            doc_id: 文档ID
        
        Returns:
            文档信息
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建索引"
                    }
            
            return self.index_manager.get_document(self.index, doc_id)
            
        except Exception as e:
            logger.error(f"获取文档失败: {e}")
            return {
                "success": False,
                "error": str(e)
            }
    
    def list_documents(self) -> Dict[str, Any]:
        """
        列出索引中的所有文档
        
        Returns:
            文档列表
        """
        try:
            # 确保索引已加载
            if self.index is None:
                if not self.load_index():
                    return {
                        "success": False,
                        "error": "索引未加载，请先构建索引",
                        "documents": []
                    }
            
            return self.index_manager.list_all_documents(self.index)
            
        except Exception as e:
            logger.error(f"列出文档失败: {e}")
            return {
                "success": False,
                "error": str(e),
                "documents": []
            }


# 便捷函数
def create_rag_system() -> RAGSystem:
    """
    创建 RAG 系统实例
    
    Returns:
        RAGSystem: RAG 系统实例
    """
    return RAGSystem()
