from concurrent.futures import ThreadPoolExecutor
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
import json

class RAGWithQueryRewriting:
    """集成查询重写的 RAG 系统"""

    def __init__(self, retriever, llm, rewrite_chain):
        self.retriever = retriever
        self.llm = llm
        self.rewrite_chain = rewrite_chain

    def invoke(self, question):
        # 第一步：重写查询
        rewritten_query = self.rewrite_chain.invoke({"query": question})
        print(f"原始问题: {question}")
        print(f"重写后:   {rewritten_query}")

        # 第二步：用重写后的查询进行检索
        docs = self.retriever.invoke(rewritten_query)

        # 第三步：用原始问题 + 检索结果生成答案
        context = "\n\n".join([doc.page_content for doc in docs])
        answer_prompt = f"""基于以下上下文回答用户问题。如果上下文中没有相关信息，请说明。

上下文：{context}

用户问题：{question}
答案："""
        answer = self.llm.invoke(answer_prompt).content

        return {
            "original_query": question,
            "rewritten_query": rewritten_query,
            "retrieved_docs": docs,
            "answer": answer
        }


class RAGWithDecomposition:
    """集成查询分解的 RAG 系统"""

    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm

    def invoke(self, question):

        decompose_prompt = PromptTemplate(
            input_variables=["question"],
            template="""你是一个问题分解助手。请将用户的复杂问题分解为 2-4 个独立的子问题，
        每个子问题应该能独立检索和回答。

        要求：
        1. 子问题之间互不依赖，可以并行检索
        2. 子问题覆盖原始问题的所有方面
        3. 每个子问题简洁明确
        4. 以 JSON 数组格式输出

        用户问题: {question}

        输出格式示例: ["子问题1", "子问题2", "子问题3"]

        子问题列表:"""
        )

        decompose_chain = decompose_prompt | self.llm | StrOutputParser()




        # 第一步：分解问题
        sub_queries = self.decompose_query(question, decompose_chain)
        print(f"分解为 {len(sub_queries)} 个子问题")

        # 第二步：并行检索并合并
        all_docs = self.parallel_retrieve_and_merge(sub_queries, self.retriever)
        print(f"合并后共 {len(all_docs)} 个文档片段")

        # 第三步：用所有上下文生成综合答案
        context = "\n\n".join([doc.page_content for doc in all_docs])
        answer_prompt = f"""基于以下上下文，全面回答用户的问题。请综合所有相关信息，给出完整、有条理的答案。上下文：{context}

用户问题：{question}
答案："""
        answer = self.llm.invoke(answer_prompt).content

        return {
            "question": question,
            "sub_queries": sub_queries,
            "doc_count": len(all_docs),
            "answer": answer
        }


    # 使用 ThreadPoolExecutor 创建线程池，并行执行检索
    # 每个子问题独立检索，互不干扰
    # 按内容哈希值去重（避免重复文档）

    def parallel_retrieve_and_merge(self, sub_queries, retriever, k_per_query=5):
        """
        对每个子问题并行检索，然后合并去重
        """
        all_docs = []
        seen_contents = set()

        def retrieve_one(query):
            return retriever.invoke(query)

        # 创建线程池，线程数等于子问题数量
        with ThreadPoolExecutor(max_workers=len(sub_queries)) as executor:
            # 提交所有检索任务
            futures = [executor.submit(retrieve_one, q) for q in sub_queries]
            ## 收集结果
            for future in futures:
                docs = future.result()
                for doc in docs:
                    # 按内容去重
                    content_hash = hash(doc.page_content)
                    if content_hash not in seen_contents:
                        seen_contents.add(content_hash)
                        all_docs.append(doc)

        return all_docs

    def decompose_query(self, question: str, decompose_chain) -> list[str]:
        """将复杂问题分解为子问题"""
        result = decompose_chain.invoke({"question": question})
        # 解析 JSON 数组
        try:
            sub_queries = json.loads(result.strip())
            return sub_queries
        except json.JSONDecodeError:
            # 兜底：按行分割
            return [line.strip() for line in result.strip().split("\n") if line.strip()]