"""
查询改写器

基于 MultiQueryRetriever 实现多路查询改写，扩大搜索范围。
"""

from typing import List, Optional, Any
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_core.language_models import BaseLanguageModel
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser


class MultiQueryTransformer:
    """
    多路查询改写器
    
    将单个查询改写成多个相关查询，用于 RAG-Fusion。
    """
    
    def __init__(
        self,
        llm: BaseLanguageModel,
        num_queries: int = 3,
        prompt_template: Optional[str] = None,
    ):
        """
        初始化查询改写器
        
        Args:
            llm: 语言模型实例
            num_queries: 生成的查询数量
            prompt_template: 提示词模板
        """
        self.llm = llm
        self.num_queries = num_queries
        
        # 默认提示词模板
        self.prompt_template = prompt_template or """
        你是一个专业的查询改写助手。你的任务是将用户的问题改写成 {num_queries} 个不同的版本。
        这些版本应该从不同的角度、使用不同的关键词来表达相同或相关的意图。
        
        原始问题: {question}
        
        请生成 {num_queries} 个不同版本的查询，每个版本一行。
        确保每个版本都是独立、完整的查询语句。
        
        生成 {num_queries} 个查询:
        """
    
    def transform_query(self, query: str) -> List[str]:
        """
        将单个查询改写成多个查询
        
        Args:
            query: 原始查询
            
        Returns:
            改写后的查询列表
        """
        prompt = PromptTemplate.from_template(self.prompt_template)
        
        chain = prompt | self.llm | StrOutputParser()
        
        response = chain.invoke({
            "question": query,
            "num_queries": self.num_queries,
        })
        
        # 解析响应，每行一个查询
        queries = [
            q.strip() 
            for q in response.strip().split('\n') 
            if q.strip()
        ]
        
        # 确保数量正确，如果不够则添加原始查询
        if len(queries) < self.num_queries:
            queries.extend([query] * (self.num_queries - len(queries)))
        elif len(queries) > self.num_queries:
            queries = queries[:self.num_queries]
        
        # 确保包含原始查询
        if query not in queries:
            queries = [query] + queries[:self.num_queries-1]
        
        return queries
    
    def create_multi_query_retriever(
        self,
        base_retriever: Any,
        include_original: bool = True,
    ) -> MultiQueryRetriever:
        """
        创建多路查询检索器
        
        Args:
            base_retriever: 基础检索器
            include_original: 是否包含原始查询
            
        Returns:
            MultiQueryRetriever 实例
        """
        retriever = MultiQueryRetriever.from_llm(
            retriever=base_retriever,
            llm=self.llm,
            include_original=include_original,
        )
        
        # 设置生成的查询数量
        retriever.llm_chain.prompt = PromptTemplate.from_template(
            "你是一个专业的查询改写助手。你的任务是将用户的问题改写成 {num_queries} 个不同的版本。\n"
            "这些版本应该从不同的角度、使用不同的关键词来表达相同或相关的意图。\n\n"
            "原始问题: {question}\n\n"
            "请生成 {num_queries} 个不同版本的查询，每个版本一行。\n"
            "确保每个版本都是独立、完整的查询语句。\n\n"
            "生成 {num_queries} 个查询:"
        )
        
        # 修改调用参数以包含 num_queries
        original_invoke = retriever.llm_chain.invoke
        
        def new_invoke(input_dict):
            input_dict["num_queries"] = self.num_queries
            return original_invoke(input_dict)
        
        retriever.llm_chain.invoke = new_invoke
        
        return retriever
    
    @classmethod
    def create_from_config(
        cls,
        llm: BaseLanguageModel,
        config: Optional[dict] = None,
    ) -> "MultiQueryTransformer":
        """
        从配置创建查询改写器
        
        Args:
            llm: 语言模型实例
            config: 配置字典
            
        Returns:
            MultiQueryTransformer 实例
        """
        config = config or {}
        return cls(
            llm=llm,
            num_queries=config.get("num_queries", 3),
            prompt_template=config.get("prompt_template", None),
        )


def create_rag_fusion_pipeline(
    base_retriever: Any,
    llm: BaseLanguageModel,
    reranker: Optional[Any] = None,
    num_queries: int = 3,
) -> Any:
    """
    创建完整的 RAG-Fusion 流水线
    
    Args:
        base_retriever: 基础检索器
        llm: 语言模型（用于查询改写）
        reranker: 重排序器（可选）
        num_queries: 查询改写数量
        
    Returns:
        检索器实例
    """
    # 创建多路查询改写器
    query_transformer = MultiQueryTransformer(
        llm=llm,
        num_queries=num_queries,
    )
    
    # 创建多路查询检索器
    multi_query_retriever = query_transformer.create_multi_query_retriever(
        base_retriever=base_retriever,
        include_original=True,
    )
    
    # 如果提供了重排序器，则应用重排序
    if reranker is not None:
        from langchain.retrievers import ContextualCompressionRetriever
        return ContextualCompressionRetriever(
            base_compressor=reranker,
            base_retriever=multi_query_retriever,
        )
    
    return multi_query_retriever