from fastapi import Body, Request
from fastapi.responses import StreamingResponse
from typing import List, Union, Dict
from loguru import logger
import importlib
import copy
import json
import os
from pathlib import Path

# from configs.model_config import (
#     llm_model_dict, LLM_MODEL, PROMPT_TEMPLATE, 
#     VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD)

from coagent.tools import (
    toLangchainTools, 
    TOOL_DICT, TOOL_SETS
)

from coagent.connector.phase import BasePhase
from coagent.connector.schema import Message
from coagent.connector.schema import Memory
from coagent.chat.utils import History, wrap_done
from coagent.llm_models.llm_config import LLMConfig, EmbedConfig
from coagent.connector.configs import PHASE_CONFIGS, AGETN_CONFIGS, CHAIN_CONFIGS

PHASE_MODULE = importlib.import_module("coagent.connector.phase")



class AgentChat:

    def __init__(
            self,
            engine_name: str = "",
            top_k: int = 1,
            stream: bool = False,
            ) -> None:
        self.top_k = top_k
        self.stream = stream
        self.chatPhase_dict: Dict[str, BasePhase] = {}

    def chat(
            self, 
            query: str = Body(..., description="用户输入", examples=["hello"]),
            phase_name: str = Body(..., description="执行场景名称", examples=["chatPhase"]),
            chain_name: str = Body(..., description="执行链的名称", examples=["chatChain"]),
            history: List[History] = Body(
                [], description="历史对话",
                examples=[[{"role": "user", "content": "我们来玩成语接龙，我先来，生龙活虎"}]]
                ),
            doc_engine_name: str = Body(..., description="知识库名称", examples=["samples"]),
            search_engine_name: str = Body(..., description="搜索引擎名称", examples=["duckduckgo"]),
            code_engine_name: str = Body(..., description="代码引擎名称", examples=["samples"]),
            top_k: int = Body(5, description="匹配向量数"),
            score_threshold: float = Body(1, description="知识库匹配相关度阈值，取值范围在0-1之间，SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右", ge=0, le=1),
            stream: bool = Body(False, description="流式输出"),
            local_doc_url: bool = Body(False, description="知识文件返回本地路径(true)或URL(false)"),
            choose_tools: List[str] = Body([], description="选择tool的集合"),
            do_search: bool = Body(False, description="是否进行搜索"),
            do_doc_retrieval: bool = Body(False, description="是否进行知识库检索"),
            do_code_retrieval: bool = Body(False, description="是否执行代码检索"),
            do_tool_retrieval: bool = Body(False, description="是否执行工具检索"),
            custom_phase_configs: dict = Body({}, description="自定义phase配置"),
            custom_chain_configs: dict = Body({}, description="自定义chain配置"),
            custom_role_configs: dict = Body({}, description="自定义role配置"),
            history_node_list: List = Body([], description="代码历史相关节点"),
            isDetailed: bool = Body(False, description="是否输出完整的agent相关内容"),
            upload_file: Union[str, Path, bytes] = "",
            kb_root_path: str = Body("", description="知识库存储路径"),
            jupyter_work_path: str = Body("", description="sandbox执行环境"),
            sandbox_server: str = Body({}, description="代码历史相关节点"),
            api_key: str = Body(os.environ.get("OPENAI_API_KEY"), description=""),
            api_base_url: str = Body(os.environ.get("API_BASE_URL"),),
            embed_model: str = Body("", description="向量模型"),
            embed_model_path: str = Body("", description="向量模型路径"),
            model_device: str = Body("", description="模型加载设备"),
            embed_engine: str = Body("", description="向量模型类型"),
            model_name: str = Body("", description="llm模型名称"),
            temperature: float = Body(0.2, description=""),
            **kargs
            ) -> Message:
        
        # update configs
        phase_configs, chain_configs, agent_configs = self.update_configs(
            custom_phase_configs, custom_chain_configs, custom_role_configs)
        params = locals()
        params.pop("self")
        embed_config: EmbedConfig = EmbedConfig(**params)
        llm_config: LLMConfig = LLMConfig(**params)

        logger.info('phase_configs={}'.format(phase_configs))
        logger.info('chain_configs={}'.format(chain_configs))
        logger.info('agent_configs={}'.format(agent_configs))
        logger.info('phase_name')
        logger.info('chain_name')

        # choose tools
        tools = toLangchainTools([TOOL_DICT[i] for i in choose_tools if i in TOOL_DICT])

        if upload_file:
            upload_file_name = upload_file if upload_file and isinstance(upload_file, str) else upload_file.name
            for _filename_idx in range(len(upload_file_name), 0, -1):
                if upload_file_name[:_filename_idx] in query:
                    query = query.replace(upload_file_name[:_filename_idx], upload_file_name)
                    break

        input_message = Message(
            role_content=query,
            role_type="user",
            role_name="human",
            input_query=query,
            origin_query=query,
            phase_name=phase_name,
            chain_name=chain_name,
            do_search=do_search,
            do_doc_retrieval=do_doc_retrieval,
            do_code_retrieval=do_code_retrieval,
            do_tool_retrieval=do_tool_retrieval,
            doc_engine_name=doc_engine_name, search_engine_name=search_engine_name,
            code_engine_name=code_engine_name,
            score_threshold=score_threshold, top_k=top_k,
            history_node_list=history_node_list,
            tools=tools
        )
        # history memory mangemant
        history = Memory(messages=[
            Message(role_name=i["role"], role_type=i["role"], role_content=i["content"]) 
            for i in history
            ])
        # start to execute
        phase_class = getattr(PHASE_MODULE, phase_configs[input_message.phase_name]["phase_type"])
        # TODO 需要把相关信息补充上去
        phase = phase_class(input_message.phase_name,
            task = input_message.task,
            base_phase_config = phase_configs,
            base_chain_config = chain_configs,
            base_role_config = agent_configs,
            phase_config = None,
            kb_root_path = kb_root_path,
            jupyter_work_path = jupyter_work_path,
            sandbox_server = sandbox_server,
            embed_config = embed_config,
            llm_config = llm_config,
            )
        output_message, local_memory = phase.step(input_message, history)

        def chat_iterator(message: Message, local_memory: Memory, isDetailed=False):
            step_content = local_memory.to_str_messages(content_key='step_content', filter_roles=["user"])
            final_content = message.role_content
            logger.debug(f"{step_content}")
            result = {
                "answer": "",
                "db_docs": [str(doc) for doc in message.db_docs],
                "search_docs": [str(doc) for doc in message.search_docs],
                "code_docs":  [str(doc) for doc in message.code_docs],
                "related_nodes": [doc.get_related_node() for idx, doc in enumerate(message.code_docs) if idx==0],
                "figures": message.figures,
                "step_content": step_content,
                "final_content": final_content,
                }
            

            related_nodes, has_nodes = [], [ ]
            for nodes in result["related_nodes"]:
                for node in nodes:
                    if node not in has_nodes:
                        related_nodes.append(node)
            result["related_nodes"] = related_nodes
            
            # logger.debug(f"{result['figures'].keys()}, isDetailed: {isDetailed}")
            message_str = step_content
            if self.stream:
                for token in message_str:
                    result["answer"] = token
                    yield json.dumps(result, ensure_ascii=False)
            else:
                for token in message_str:
                    result["answer"] += token
                yield json.dumps(result, ensure_ascii=False)
            
        return StreamingResponse(chat_iterator(output_message, local_memory, isDetailed), media_type="text/event-stream")


    def achat(
            self, 
            query: str = Body(..., description="用户输入", examples=["hello"]),
            phase_name: str = Body(..., description="执行场景名称", examples=["chatPhase"]),
            chain_name: str = Body(..., description="执行链的名称", examples=["chatChain"]),
            history: List[History] = Body(
                [], description="历史对话",
                examples=[[{"role": "user", "content": "我们来玩成语接龙，我先来，生龙活虎"}]]
                ),
            doc_engine_name: str = Body(..., description="知识库名称", examples=["samples"]),
            search_engine_name: str = Body(..., description="搜索引擎名称", examples=["duckduckgo"]),
            code_engine_name: str = Body(..., description="代码引擎名称", examples=["samples"]),
            cb_search_type: str = Body(..., description="代码查询模式", examples=["tag"]),
            top_k: int = Body(5, description="匹配向量数"),
            score_threshold: float = Body(1, description="知识库匹配相关度阈值，取值范围在0-1之间，SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右", ge=0, le=1),
            stream: bool = Body(False, description="流式输出"),
            local_doc_url: bool = Body(False, description="知识文件返回本地路径(true)或URL(false)"),
            choose_tools: List[str] = Body([], description="选择tool的集合"),
            do_search: bool = Body(False, description="是否进行搜索"),
            do_doc_retrieval: bool = Body(False, description="是否进行知识库检索"),
            do_code_retrieval: bool = Body(False, description="是否执行代码检索"),
            do_tool_retrieval: bool = Body(False, description="是否执行工具检索"),
            custom_phase_configs: dict = Body({}, description="自定义phase配置"),
            custom_chain_configs: dict = Body({}, description="自定义chain配置"),
            custom_role_configs: dict = Body({}, description="自定义role配置"),
            history_node_list: List = Body([], description="代码历史相关节点"),
            isDetailed: bool = Body(False, description="是否输出完整的agent相关内容"),
            upload_file: Union[str, Path, bytes] = "",
            kb_root_path: str = Body("", description="知识库存储路径"),
            jupyter_work_path: str = Body("", description="sandbox执行环境"),
            sandbox_server: str = Body({}, description="代码历史相关节点"),
            api_key: str = Body(os.environ["OPENAI_API_KEY"], description=""),
            api_base_url: str = Body(os.environ.get("API_BASE_URL"),),
            embed_model: str = Body("", description="向量模型"),
            embed_model_path: str = Body("", description="向量模型路径"),
            model_device: str = Body("", description="模型加载设备"),
            embed_engine: str = Body("", description="向量模型类型"),
            model_name: str = Body("", description="llm模型名称"),
            temperature: float = Body(0.2, description=""),
            **kargs
            ) -> Message:
        
        # update configs
        phase_configs, chain_configs, agent_configs = self.update_configs(
            custom_phase_configs, custom_chain_configs, custom_role_configs)
        
        # 
        params = locals()
        params.pop("self")
        embed_config: EmbedConfig = EmbedConfig(**params)
        llm_config: LLMConfig = LLMConfig(**params)

        # choose tools
        tools = toLangchainTools([TOOL_DICT[i] for i in choose_tools if i in TOOL_DICT])

        if upload_file:
            upload_file_name = upload_file if upload_file and isinstance(upload_file, str) else upload_file.name
            for _filename_idx in range(len(upload_file_name), 0, -1):
                if upload_file_name[:_filename_idx] in query:
                    query = query.replace(upload_file_name[:_filename_idx], upload_file_name)
                    break

        input_message = Message(
            role_content=query,
            role_type="user",
            role_name="human",
            input_query=query,
            origin_query=query,
            phase_name=phase_name,
            chain_name=chain_name,
            do_search=do_search,
            do_doc_retrieval=do_doc_retrieval,
            do_code_retrieval=do_code_retrieval,
            do_tool_retrieval=do_tool_retrieval,
            doc_engine_name=doc_engine_name, 
            search_engine_name=search_engine_name,
            code_engine_name=code_engine_name,
            cb_search_type=cb_search_type,
            score_threshold=score_threshold, top_k=top_k,
            history_node_list=history_node_list,
            tools=tools
        )
        # history memory mangemant
        history = Memory(messages=[
            Message(role_name=i["role"], role_type=i["role"], role_content=i["content"]) 
            for i in history
            ])
        # start to execute
        if  phase_configs[input_message.phase_name]["phase_type"] not in self.chatPhase_dict:
            phase_class = getattr(PHASE_MODULE, phase_configs[input_message.phase_name]["phase_type"])
            phase = phase_class(input_message.phase_name,
                task = input_message.task,
                base_phase_config = phase_configs,
                base_chain_config = chain_configs,
                base_role_config = agent_configs,
                phase_config = None,
                kb_root_path = kb_root_path,
                jupyter_work_path = jupyter_work_path,
                sandbox_server = sandbox_server,
                embed_config = embed_config,
                llm_config = llm_config,
                )
            self.chatPhase_dict[phase_configs[input_message.phase_name]["phase_type"]] = phase
        else:
            phase = self.chatPhase_dict[phase_configs[input_message.phase_name]["phase_type"]]
        
        def chat_iterator(message: Message, local_memory: Memory, isDetailed=False):
            step_content = local_memory.to_str_messages(content_key='step_content', filter_roles=["human"])
            step_content = "\n\n".join([f"{v}" for parsed_output in local_memory.get_parserd_output_list()[1:] for k, v in parsed_output.items() if k not in ["Action Status"]])
            final_content = message.role_content
            result = {
                "answer": "",
                "db_docs": [str(doc) for doc in message.db_docs],
                "search_docs": [str(doc) for doc in message.search_docs],
                "code_docs":  [str(doc) for doc in message.code_docs],
                "related_nodes": [doc.get_related_node() for idx, doc in enumerate(message.code_docs) if idx==0],
                "figures": message.figures,
                "step_content": step_content or final_content,
                "final_content": final_content,
                }
            
            related_nodes, has_nodes = [], [ ]
            for nodes in result["related_nodes"]:
                for node in nodes:
                    if node not in has_nodes:
                        related_nodes.append(node)
            result["related_nodes"] = related_nodes
            
            # logger.debug(f"{result['figures'].keys()}, isDetailed: {isDetailed}")
            message_str = step_content
            if self.stream:
                for token in message_str:
                    result["answer"] = token
                    yield json.dumps(result, ensure_ascii=False)
            else:
                for token in message_str:
                    result["answer"] += token
                yield json.dumps(result, ensure_ascii=False)
            

        for output_message, local_memory in phase.astep(input_message, history):

            # logger.debug(f"output_message: {output_message}")
            # output_message = Message(**output_message)
            # local_memory = Memory(**local_memory)
            for result in chat_iterator(output_message, local_memory, isDetailed):
                yield result


    def _chat(self, ):
        pass

    def update_configs(self, custom_phase_configs, custom_chain_configs, custom_role_configs):
        '''update phase/chain/agent configs'''
        phase_configs = copy.deepcopy(PHASE_CONFIGS)
        phase_configs.update(custom_phase_configs)
        chain_configs = copy.deepcopy(CHAIN_CONFIGS)
        chain_configs.update(custom_chain_configs)
        agent_configs = copy.deepcopy(AGETN_CONFIGS)
        agent_configs.update(custom_role_configs)
        # phase_configs = load_phase_configs(new_phase_configs)
        # chian_configs = load_chain_configs(new_chain_configs)
        # agent_configs = load_role_configs(new_agent_configs)
        return phase_configs, chain_configs, agent_configs