"""
查询引擎模块
负责处理查询和检索操作，支持LlamaIndex后处理器和元数据过滤
"""
from typing import Optional, List, Dict, Any
import asyncio
from loguru import logger

from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.vector_stores import MetadataFilters, MetadataFilter, FilterOperator

from src.config import get_config_manager, ConfigManager
from src.storage import BM25StoreManager
from src.utils.validators import validate_bm25_retriever, apply_nest_asyncio
from src.utils.decorators import handle_query_errors, handle_retrieval_errors
from src.utils.log_formatter import log_retrieval_result, log_fallback


class QueryEngineManager:
    """查询引擎管理器类（增强版，支持后处理器和元数据过滤）"""
    
    def __init__(self, config: Optional[ConfigManager] = None):
        """
        初始化查询引擎管理器
        
        Args:
            config: 配置实例，如果为None则使用全局配置
        """
        self.config = config or get_config_manager()
        self.bm25_store_manager = BM25StoreManager(config=self.config)
    
    def _create_postprocessors(self) -> List:
        """
        创建后处理器链
        基于LlamaIndex官方组件
        
        Returns:
            后处理器列表
        """
        postprocessors = []
        retrieval_config = self.config.retrieval
        
        # 1. 相似度过滤 - SimilarityPostprocessor
        if retrieval_config.use_similarity_filter:
            try:
                from llama_index.core.postprocessor import SimilarityPostprocessor
                postprocessors.append(
                    SimilarityPostprocessor(
                        similarity_cutoff=retrieval_config.similarity_cutoff
                    )
                )
                logger.debug(f"✅ 已启用相似度过滤，阈值: {retrieval_config.similarity_cutoff}")
            except ImportError as e:
                logger.warning(f"⚠️ SimilarityPostprocessor 导入失败: {e}")
        
        # 2. 多样性优化 - SentenceEmbeddingOptimizer
        if retrieval_config.use_optimizer:
            try:
                from llama_index.core.postprocessor import SentenceEmbeddingOptimizer
                # 需要embedding model
                from src.di import get_container
                container = get_container()
                embed_model = container.resolve("embed_model")
                
                postprocessors.append(
                    SentenceEmbeddingOptimizer(
                        embed_model=embed_model,
                        percentile_cutoff=retrieval_config.optimizer_percentile,
                        threshold_cutoff=retrieval_config.optimizer_threshold
                    )
                )
                logger.debug(f"✅ 已启用多样性优化器")
            except ImportError as e:
                logger.warning(f"⚠️ SentenceEmbeddingOptimizer 导入失败: {e}")
            except Exception as e:
                logger.warning(f"⚠️ 多样性优化器创建失败: {e}")
        
        # 3. 重排序 - Reranker
        if retrieval_config.use_reranker:
            reranker = self._create_reranker()
            if reranker:
                postprocessors.append(reranker)
        
        return postprocessors
    
    def _create_reranker(self):
        """
        创建重排序器
        支持多种LlamaIndex官方Reranker插件
        
        Returns:
            Reranker实例或None
        """
        retrieval_config = self.config.retrieval
        reranker_type = retrieval_config.reranker_type
        
        try:
            if reranker_type == "cohere":
                # Cohere Rerank - 商业级质量
                from llama_index.postprocessor.cohere_rerank import CohereRerank
                api_key = retrieval_config.get_cohere_api_key()
                if not api_key:
                    logger.warning("⚠️ Cohere API Key 未配置，跳过Cohere Rerank")
                    return None
                reranker = CohereRerank(
                    api_key=api_key,
                    top_n=retrieval_config.reranker_top_n
                )
                logger.debug(f"✅ 已启用 Cohere Rerank，top_n: {retrieval_config.reranker_top_n}")
                return reranker
            
            elif reranker_type == "sentence-transformer":
                # SentenceTransformer Rerank - 免费本地
                from llama_index.core.postprocessor import SentenceTransformerRerank
                reranker = SentenceTransformerRerank(
                    model=retrieval_config.reranker_model,
                    top_n=retrieval_config.reranker_top_n
                )
                logger.debug(f"✅ 已启用 SentenceTransformer Rerank，模型: {retrieval_config.reranker_model}")
                return reranker
            
            elif reranker_type == "flag-embedding":
                # Flag Embedding Reranker - 中文优化
                from llama_index.postprocessor.flag_embedding_reranker import FlagEmbeddingReranker
                reranker = FlagEmbeddingReranker(
                    top_n=retrieval_config.reranker_top_n,
                    model="BAAI/bge-reranker-large"
                )
                logger.debug(f"✅ 已启用 Flag Embedding Rerank")
                return reranker
            
            else:
                logger.warning(f"⚠️ 不支持的Reranker类型: {reranker_type}")
                return None
        
        except ImportError as e:
            logger.warning(f"⚠️ Reranker插件未安装 ({reranker_type}): {e}")
            logger.info(f"💡 安装提示: pip install llama-index-postprocessor-{reranker_type}")
            return None
        except Exception as e:
            logger.error(f"❌ Reranker创建失败: {e}")
            return None
    
    def _create_metadata_filters(self, metadata_filters: Dict) -> Optional[MetadataFilters]:
        """
        创建元数据过滤器
        基于LlamaIndex原生MetadataFilters
        
        Args:
            metadata_filters: 过滤器定义字典
                格式: {"key": {"operator": "EQ", "value": "xxx"}}
        
        Returns:
            MetadataFilters实例或None
        """
        if not metadata_filters or not self.config.retrieval.enable_metadata_filters:
            return None
        
        try:
            filter_list = []
            for key, filter_def in metadata_filters.items():
                operator_str = filter_def.get("operator", "EQ").upper()
                operator = FilterOperator[operator_str]
                value = filter_def.get("value")
                
                filter_list.append(
                    MetadataFilter(key=key, value=value, operator=operator)
                )
            
            filters = MetadataFilters(filters=filter_list)
            logger.debug(f"✅ 已创建元数据过滤器: {len(filter_list)} 个条件")
            return filters
        
        except Exception as e:
            logger.error(f"❌ 元数据过滤器创建失败: {e}")
            return None
    
    def create_query_engine(
        self,
        index: VectorStoreIndex,
        mode: Optional[str] = None,
        top_k: Optional[int] = None,
        metadata_filters: Optional[Dict] = None,
        similarity_cutoff: Optional[float] = None,
        use_reranker: Optional[bool] = None
    ):
        """
        创建查询引擎（增强版）
        
        Args:
            index: 向量索引实例
            mode: 检索模式
            top_k: 检索数量
            metadata_filters: 元数据过滤条件
            similarity_cutoff: 相似度阈值（覆盖配置）
            use_reranker: 是否使用重排序（覆盖配置）
            
        Returns:
            查询引擎实例
        """
        retrieval_config = self.config.retrieval
        retrieval_mode = mode or retrieval_config.mode
        similarity_top_k = top_k or retrieval_config.top_k
        
        # 创建后处理器（如果启用）
        postprocessors = self._create_postprocessors()
        
        # 动态覆盖相似度阈值
        if similarity_cutoff is not None:
            from llama_index.core.postprocessor import SimilarityPostprocessor
            # 移除已有的相似度过滤器
            postprocessors = [p for p in postprocessors if not isinstance(p, SimilarityPostprocessor)]
            # 添加新的阈值
            postprocessors.insert(0, SimilarityPostprocessor(similarity_cutoff=similarity_cutoff))
        
        # 动态控制Reranker
        if use_reranker is False:
            # 移除所有Reranker
            from llama_index.core.postprocessor import BaseNodePostprocessor
            postprocessors = [p for p in postprocessors 
                            if not (hasattr(p, '__class__') and 'rerank' in p.__class__.__name__.lower())]
        
        # 创建元数据过滤器
        filters = self._create_metadata_filters(metadata_filters)
        
        if retrieval_mode == "vector":
            # 纯向量检索
            return index.as_query_engine(
                similarity_top_k=similarity_top_k,
                filters=filters,
                node_postprocessors=postprocessors,
                response_mode="compact",
            )
        elif retrieval_mode == "bm25":
            # 纯BM25检索
            bm25_retriever = self.bm25_store_manager.get_bm25_retriever()
            validate_bm25_retriever(retrieval_config.use_bm25, bm25_retriever, raise_error=True)
            
            return RetrieverQueryEngine(
                retriever=bm25_retriever,
                node_postprocessors=postprocessors,
            )
        elif retrieval_mode == "hybrid":
            # 混合检索
            return self._create_hybrid_query_engine(
                index, similarity_top_k, filters, postprocessors
            )
        else:
            raise ValueError(f"不支持的检索模式: {retrieval_mode}")
    
    def _create_hybrid_query_engine(
        self, 
        index: VectorStoreIndex, 
        similarity_top_k: int,
        filters: Optional[MetadataFilters] = None,
        postprocessors: Optional[List] = None
    ):
        """
        创建混合查询引擎
        
        Args:
            index: 向量索引实例
            similarity_top_k: 检索数量
            filters: 元数据过滤器
            postprocessors: 后处理器列表
            
        Returns:
            混合查询引擎实例
        """
        bm25_retriever = self.bm25_store_manager.get_bm25_retriever()
        retrieval_config = self.config.retrieval
        
        # 检查BM25是否可用，如果不可用则回退
        if not validate_bm25_retriever(retrieval_config.use_bm25, bm25_retriever, raise_error=False):
            log_fallback("hybrid", "vector", "BM25检索器未启用")
            return index.as_query_engine(
                similarity_top_k=similarity_top_k,
                filters=filters,
                node_postprocessors=postprocessors or [],
                response_mode="compact",
            )
        
        # 应用nest_asyncio以支持异步
        apply_nest_asyncio()
        
        # 使用Qdrant时禁用异步模式
        use_async_mode = not self.config.vector_store.use_qdrant
        
        retriever = QueryFusionRetriever(
            retrievers=[
                index.as_retriever(similarity_top_k=similarity_top_k, filters=filters),
                bm25_retriever,
            ],
            num_queries=1,
            use_async=use_async_mode,
        )
        
        return RetrieverQueryEngine(
            retriever=retriever,
            node_postprocessors=postprocessors or [],
        )
    
    def create_retriever(
        self,
        index: VectorStoreIndex,
        mode: Optional[str] = None,
        top_k: Optional[int] = None,
        metadata_filters: Optional[Dict] = None
    ):
        """
        创建检索器（仅检索，不使用LLM）
        
        Args:
            index: 向量索引实例
            mode: 检索模式
            top_k: 检索数量
            metadata_filters: 元数据过滤条件
            
        Returns:
            检索器实例
        """
        retrieval_config = self.config.retrieval
        retrieval_mode = mode if mode is not None else "hybrid"
        similarity_top_k = top_k or retrieval_config.top_k
        
        # 创建元数据过滤器
        filters = self._create_metadata_filters(metadata_filters)
        
        if retrieval_mode == "vector":
            # 纯向量检索
            return index.as_retriever(
                similarity_top_k=similarity_top_k,
                filters=filters
            )
        elif retrieval_mode == "bm25":
            # 纯BM25检索
            bm25_retriever = self.bm25_store_manager.get_bm25_retriever()
            if not retrieval_config.use_bm25 or bm25_retriever is None:
                raise ValueError("BM25检索器未启用或未加载")
            return bm25_retriever
        elif retrieval_mode == "hybrid":
            # 混合检索
            return self._create_hybrid_retriever(index, similarity_top_k, filters)
        else:
            raise ValueError(f"不支持的检索模式: {retrieval_mode}")
    
    def _create_hybrid_retriever(
        self, 
        index: VectorStoreIndex, 
        similarity_top_k: int,
        filters: Optional[MetadataFilters] = None
    ):
        """
        创建混合检索器
        
        Args:
            index: 向量索引实例
            similarity_top_k: 检索数量
            filters: 元数据过滤器
            
        Returns:
            混合检索器实例
        """
        bm25_retriever = self.bm25_store_manager.get_bm25_retriever()
        retrieval_config = self.config.retrieval
        
        if not retrieval_config.use_bm25 or bm25_retriever is None:
            logger.warning("BM25检索器未启用，回退到纯向量检索")
            return index.as_retriever(
                similarity_top_k=similarity_top_k,
                filters=filters
            )
        
        # 应用nest_asyncio以支持异步
        try:
            import nest_asyncio
            nest_asyncio.apply()
        except RuntimeError:
            logger.debug("nest_asyncio 已应用或事件循环已在运行")
        
        # 使用Qdrant时禁用异步模式
        use_async_mode = not self.config.vector_store.use_qdrant
        
        return QueryFusionRetriever(
            retrievers=[
                index.as_retriever(similarity_top_k=similarity_top_k, filters=filters),
                bm25_retriever,
            ],
            num_queries=1,
            use_async=use_async_mode,
        )
    
    @handle_query_errors
    def execute_query(self, query_engine, question: str, mode: str) -> dict:
        """
        执行查询
        
        Args:
            query_engine: 查询引擎实例
            question: 问题
            mode: 检索模式
            
        Returns:
            dict: 查询结果
        """
        response = query_engine.query(question)
        
        # 提取源文档信息
        sources = []
        if hasattr(response, 'source_nodes'):
            for node in response.source_nodes:
                sources.append({
                    "text": node.node.text[:200] + "..." if len(node.node.text) > 200 else node.node.text,
                    "score": node.score if hasattr(node, 'score') else None,
                    "metadata": node.node.metadata,
                })
        
        log_retrieval_result(mode, len(sources), "查询")
        
        return {
            "success": True,
            "answer": str(response),
            "sources": sources,
            "mode": mode,
            "error": None
        }
    
    def execute_retrieve(self, retriever, question: str, mode: str) -> dict:
        """
        执行检索
        
        Args:
            retriever: 检索器实例
            question: 问题
            mode: 检索模式
            
        Returns:
            dict: 检索结果
        """
        try:
            nodes = retriever.retrieve(question)
            
            # 提取文档信息
            results = []
            for idx, node in enumerate(nodes, 1):
                metadata = node.node.metadata or {}
                file_path = metadata.get("file_path", "")
                file_name = metadata.get("file_name", "")
                
                if not file_name and file_path:
                    from pathlib import Path
                    file_name = Path(file_path).name
                
                score = node.score if hasattr(node, 'score') else None
                
                result_item = {
                    "rank": idx,
                    "file_path": file_path,
                    "file_name": file_name,
                    "score": score,
                    "text": node.node.text,
                    "metadata": metadata,
                }
                results.append(result_item)
            
            logger.info(f"✅ 检索完成 - 检索模式: {mode}, 检索到 {len(results)} 个相关文档片段")
            
            return {
                "success": True,
                "results": results,
                "mode": mode,
                "count": len(results),
                "query": question,
                "error": None
            }
            
        except Exception as e:
            logger.error(f"检索失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e),
                "results": [],
                "mode": mode
            }
    
    async def batch_query_async(
        self, 
        query_engine, 
        questions: List[str]
    ) -> List[Dict[str, Any]]:
        """
        异步批量查询
        
        Args:
            query_engine: 查询引擎实例
            questions: 问题列表
            
        Returns:
            查询结果列表
        """
        try:
            tasks = [query_engine.aquery(q) for q in questions]
            responses = await asyncio.gather(*tasks)
            
            results = []
            for idx, (question, response) in enumerate(zip(questions, responses), 1):
                sources = []
                if hasattr(response, 'source_nodes'):
                    for node in response.source_nodes:
                        sources.append({
                            "text": node.node.text[:200] + "..." if len(node.node.text) > 200 else node.node.text,
                            "score": node.score if hasattr(node, 'score') else None,
                            "metadata": node.node.metadata,
                        })
                
                results.append({
                    "index": idx,
                    "question": question,
                    "answer": str(response),
                    "sources": sources,
                    "success": True
                })
            
            logger.info(f"✅ 批量查询完成 - 共处理 {len(questions)} 个问题")
            return results
        
        except Exception as e:
            logger.error(f"❌ 批量查询失败: {e}")
            return [{
                "index": i,
                "question": q,
                "answer": "",
                "sources": [],
                "success": False,
                "error": str(e)
            } for i, q in enumerate(questions, 1)]
    
    def batch_query(
        self, 
        query_engine, 
        questions: List[str]
    ) -> List[Dict[str, Any]]:
        """
        批量查询（同步接口）
        
        Args:
            query_engine: 查询引擎实例
            questions: 问题列表
            
        Returns:
            查询结果列表
        """
        return asyncio.run(self.batch_query_async(query_engine, questions))
