import sys
import os
import json
from typing import List, Dict, Any
import logging
import time
import asyncio
import httpx
sys.path.append(os.getcwd())

from rag2_0.dify.dify_client.client import DifyClient, KnowledgeBaseClient
from rag2_0.tool.ModelTool import XinferenceReRankerModel
class DifyQueryRetrieval:

    def __init__(self, dify_dataset_key: str, dify_base_url: str):
        self._dify_dataset_key = dify_dataset_key
        self._dify_base_url = dify_base_url
        self._datasets_list = self.get_datasets_list()
        
    def get_datasets_list(self) -> Dict[str, str]:
        client = KnowledgeBaseClient(api_key=self._dify_dataset_key, base_url=self._dify_base_url)
        datasets = client.list_datasets(page_size=50)
        datasets_json = datasets.json()
        return {dataset["name"]:dataset for dataset in datasets_json["data"]}

    def retrieve_by_dataset(self, query: str, dataset_name: str) -> Dict[str, Any]:
        try:
            dataset_id = self._datasets_list[dataset_name]["id"]
            retrieval_model = self._datasets_list[dataset_name]["retrieval_model_dict"]
            knowledge_base_client = KnowledgeBaseClient(api_key=self._dify_dataset_key, base_url=self._dify_base_url, dataset_id=dataset_id)
            
            documents = knowledge_base_client.retrieve(query, retrieval_model=retrieval_model, timeout=300)
            retrieved_documents = documents.json().get("records", [])
            
            # 添加数据集信息
            for retrieved_document in retrieved_documents:
                retrieved_document["dataset_id"] = dataset_id
                retrieved_document["dataset_name"] = dataset_name
            
            # 返回包含查询和文档的字典
            return {
                "query": query,
                "dataset_name": dataset_name,
                "documents": retrieved_documents
            }
        except Exception as e:
            logging.error(f"检索数据集 {dataset_name} 时出错: {str(e)}", exc_info=True)
            return {
                "query": query,
                "dataset_name": dataset_name,
                "documents": []
            }

    async def retrieve_by_dataset_async(self, query: str, dataset_name: str) -> Dict[str, Any]:
        """
        异步版本的retrieve_by_dataset方法
        
        Args:
            query: 查询字符串
            dataset_name: 数据集名称
            
        Returns:
            包含查询、数据集名称和检索到的文档的字典
        """
        try:
            # 使用asyncio.to_thread包装同步方法
            return await asyncio.to_thread(
                self.retrieve_by_dataset,
                query,
                dataset_name
            )
        except Exception as e:
            logging.error(f"异步检索数据集 {dataset_name} 时出错: {str(e)}", exc_info=True)
            return {
                "query": query,
                "dataset_name": dataset_name,
                "documents": []
            }

    async def retrieve_api_async(self, original_query: str, query_list: List[str], data_set_list: List[str], query_expand_dict: dict, top_k: int = 5)->Dict[str, Any]:
        """
        异步版本的retrieve_api方法，使用asyncio代替线程池
        
        Args:
            original_query: 原始查询
            query_list: 查询列表
            data_set_list: 数据集列表
            query_expand_dict: 查询扩展字典，包含不同类型的查询
            top_k: 返回的文档数量
            
        Returns:
            包含检索结果和查询命中统计的字典
        """
        all_documents = []
        query_document_mapping = {}  # 用于存储查询和文档的映射关系，键为查询，值为文档列表
        # 记录开始时间
        time_start = time.time()
        
        # 创建异步任务列表
        tasks = []
        for query in query_list:
            for dataset in data_set_list:
                if dataset not in list(self._datasets_list.keys()):
                    logging.error(f"dataset {dataset} not in datasets_list")
                    continue
                
                # 创建异步任务
                task = self.retrieve_by_dataset_async(query, dataset)
                tasks.append(task)
        
        # 并发执行所有异步任务
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果
        for result in results:
            if isinstance(result, Exception):
                logging.error(f"异步检索过程中发生错误: {str(result)}", exc_info=True)
            else:
                # 添加查询和文档的映射关系
                query = result["query"]
                if query not in query_document_mapping:
                    query_document_mapping[query] = []
                
                # 将文档添加到对应查询的文档列表中
                query_document_mapping[query].extend([item['segment']['id'] for item in result["documents"]])
                
                # 将文档添加到总文档列表
                all_documents.extend(result["documents"])
        
        time_end = time.time()
        logging.info(f"异步检索耗时: {time_end - time_start:.2f}秒")
        
        # 根据segment_id对文档进行去重
        unique_documents = {}
        for document in all_documents:
            segment_id = document['segment']['id']
            if segment_id not in unique_documents:
                unique_documents[segment_id] = document
        
        # 将去重后的文档转换为列表
        deduplicated_documents = list(unique_documents.values())
        
        if len(deduplicated_documents) == 0:
            return {
                "documents": [],
                "query_hit_stats": {}
            }
            
        # 对所有检索出来的文档进行重排序
        time_start = time.time()
        processed_documents = await self.data_post_processor_async(original_query, deduplicated_documents, top_k)
        time_end = time.time()
        logging.info(f"异步检索后重排序耗时: {time_end - time_start:.2f}秒")
        
        # 统计不同类型查询命中的文档
        query_hit_stats = {}
        
        # 获取重排序后的文档ID列表
        reranked_doc_ids = [doc["metadata"]["segment_id"] for doc in processed_documents]
        reranked_doc_titles = [doc["title"].split("/")[-1] for doc in processed_documents]
        
        # 解析query_expand_dict（如果是字符串）
        if isinstance(query_expand_dict, str):
            try:
                query_expand_dict = json.loads(query_expand_dict)
            except Exception as e:
                logging.error(f"解析query_expand_dict失败: {str(e)}")
                query_expand_dict = {}
        
        # 统计各类型查询命中的文档
        for query_type, queries in query_expand_dict.items():
            if not isinstance(queries, list):
                queries = [queries]
                
            # 初始化该查询类型命中的文档列表
            query_hit_stats[query_type] = []
            
            # 合并所有该类型查询命中的文档
            hit_doc_ids_set = set()
            for query in queries:
                if query in query_document_mapping:
                    hit_doc_ids_set.update(set(query_document_mapping[query]))
            
            # 找出在重排序结果中的文档
            hit_doc_titles_set = set()  # 用于去重
            for i, doc_id in enumerate(reranked_doc_ids):
                if doc_id in hit_doc_ids_set:
                    doc_title = reranked_doc_titles[i]
                    if doc_title not in hit_doc_titles_set:  # 确保不添加重复的标题
                        hit_doc_titles_set.add(doc_title)
                        query_hit_stats[query_type].append(doc_title)

        
        return {
            "documents": processed_documents,
            "query_hit_stats": query_hit_stats
        }

    async def data_post_processor_async(self, query: str, all_documents: List[Dict[str, Any]], top_k: int = 5) -> List[Dict[str, Any]]:
        """
        异步版本的data_post_processor方法
        
        Args:
            query: 查询字符串
            all_documents: 待处理的文档列表
            
        Returns:
            处理后的文档列表
        """
        reranker_model = XinferenceReRankerModel()
        documents = [document['segment']['content'] for document in all_documents]
        # 使用异步重排序方法
        reranked_documents = await reranker_model.rerank_async(query, documents, top_k=top_k)
        new_all_documents = []

        def to_dify_document_format(document: dict)->dict:
            return {
                "metadata": {
                    "_source": "knowledge",
                    "dataset_id": document["dataset_id"],
                    "dataset_name": document["dataset_name"],
                    "document_id": document['segment']['document_id'],
                    "document_name": document["segment"]["document"]["name"],
                    "data_source_type": document["segment"]["document"]["data_source_type"],
                    "segment_id": document["segment"]["id"],
                    "retriever_from": "api",
                    "score": document.get("score", 0),
                    "segment_hit_count": document.get("segment", {}).get("hit_count", 0),
                    "segment_word_count": document.get("segment", {}).get("word_count", 0),
                    "segment_position": document.get("segment", {}).get("position", 0),
                    "segment_index_node_hash": document.get("segment", {}).get("index_node_hash", ""),
                    "doc_metadata": document.get("segment", {}).get("document", {}).get("doc_metadata", None),
                    "position": document["segment"].get("position", 0)
                },
                "title": document["segment"]["document"]["name"],
                "content": document["segment"]["content"]
            }

        for reranked_document in reranked_documents:
            cur_doc_info = all_documents[reranked_document["index"]]
            cur_doc_info["score"] = reranked_document["score"]
            new_all_documents.append(to_dify_document_format(cur_doc_info))
        return new_all_documents

    
if __name__ == "__main__":
    dify_query_retrieval = DifyQueryRetrieval(api_key="dataset-skLjmPVonjHo119OWNf3kAmY", base_url="http://10.1.16.39/v1")
    
    # 测试异步API
    async def test_async_api():
        datasets = await dify_query_retrieval.retrieve_api_async(
            "电力建设计价通软件如何批量修改设备价格？", 
            ["电力建设计价通软件如何批量修改设备价格？"], 
            ["电力建设计价通（2018）软件知识(new)"]
        )
        print("异步API测试结果:")
        print(json.dumps(datasets, ensure_ascii=False, indent=2))
    
    # 如果需要测试异步API，取消下面的注释
    import asyncio
    asyncio.run(test_async_api())