import json
import os
import time
import operator
import difflib

from typing import TypedDict, List, Annotated
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage

from pydantic import BaseModel, Field

from src.multi_agent.util.tools import formatted_json
from src.multi_agent.util.web_search import do_search
from src.multi_agent.config import LLM
from src.multi_agent.weather_agent import weather_chain

llm = LLM

MEMBERS = {'天气查询模块': '输入一个城市和日期，查询当日的天气状况',
           '网络查询模块': '在网上搜索指定的内容',
           '诗歌创作模块': '根据给定的要求，进行诗歌创作',
           '结束节点': '完成所有任务或者任务执行失败的情况下进入'}


class AgentState(TypedDict):
    user_query: str
    messages: Annotated[List[str], operator.add]
    plan: str
    members: dict
    next_agent: str
    next_plan: str


def planner(state: AgentState):
    user_query = state.get('user_query', '')
    planner_prompt = ChatPromptTemplate.from_messages([
        ('system', '''你是一个团队的管理者，你负责根据用户输入的要求，制定一个计划来完成。你的团队成员和能力如下：
1. 天气查询模块：输入一个城市和日期，查询当日的天气状况。
2. 网络查询模块：输入一个需要查询的内容，在网上搜索指定的内容。
3. 诗歌创作模块：根据给定的要求，进行诗歌创作
请你根据用户的需求，制定一个处理计划'''),
        ('user', f'当前的用户需求是:'),
        MessagesPlaceholder(variable_name="query"),
    ])
    planner_chain = (planner_prompt | llm)

    llm_resp = planner_chain.invoke({'query': [('user', user_query)]})
    return {'plan': llm_resp.content}


def decision_node(state: AgentState):
    members = state.get('members', {})
    member_keys = list(members.keys())
    if isinstance(llm, ChatOpenAI):
        class NextStep(BaseModel):
            next_agent: str = Field(description='下一步应该交给哪个模块来完成')
            action: str = Field(description='下一步应该执行的行动计划是什么。如果已经执行所有的任务，则填入执行结果')

        decision_prompt = ChatPromptTemplate.from_messages([
            ('system', f'你负责针对一个给定的任务计划，以及当前的任务执行状态，制定下一步应该由谁来执行什么内容。'
                       f'可以在这些模块中进行选择：{", ".join(member_keys)}。'),
            ('user', f'当前的计划是:'),
            MessagesPlaceholder(variable_name="plan"),
            ('user', '以下是各个模块之间的消息记录'),
            MessagesPlaceholder(variable_name="messages")
        ])
        decision_chain = (decision_prompt | llm.with_structured_output(NextStep))
        llm_resp = decision_chain.invoke({'plan': [('user', state.get('plan', ''))],
                                          'messages': state.get('messages', [])})
        next_agent = llm_resp.next_agent
        action = llm_resp.action
    else:
        decision_prompt = ChatPromptTemplate.from_messages([
            ('system',
             f'你负责针对一个给定的任务计划，以及当前的任务执行状态，制定下一步应该由谁来执行什么内容。'
             f'可以在这些模块中进行选择：{", ".join(member_keys)}。你的输出必须只能是一个json字符串，'
             f'必须包含next_agent和action两个字段。其中next_agent是下一步应该交给哪个模块来完成，'
             f'action是下一步应该执行的行动计划是什么，如果已经执行完成所有的任务，则在action中填入执行结果'
             f'你只需要输出这个json即可，不需要给出任何解释'),
            ('user', f'当前的计划是:'),
            MessagesPlaceholder(variable_name="plan"),
            ('user', '以下是各个模块之间的消息记录'),
            MessagesPlaceholder(variable_name="messages")
        ])
        decision_chain = (decision_prompt | llm)
        llm_resp = decision_chain.invoke(
            {'plan': [('user', state.get('plan', ''))], 'messages': state.get('messages', [])})
        llm_resp_str = formatted_json(llm_resp.content)
        resp_dict = json.loads(llm_resp_str)
        next_agent = resp_dict.get('next_agent')
        action = resp_dict.get('action')
        if not (next_agent and action):
            raise ValueError('模型返回异常')
    return {'messages': [HumanMessage(content=f'接下来由{next_agent}来执行{action}')], 'next_agent': next_agent, 'next_plan': action}


def router_node(state: AgentState):
    members = state.get('members', {})
    member_keys = list(members.keys())
    best_match = difflib.get_close_matches(state.get('next_agent', ''), member_keys, n=1, cutoff=0)
    if not best_match:
        raise ValueError(f'Select a wrong agent: {state.get("next_agent")}')
    return {'next': best_match}


def weather_node(state: AgentState):
    # result = '上海的天气是25℃，下雨'
    agent_resp = weather_chain.invoke({'next_plan': state.get('next_plan')})
    result = agent_resp.get('result')
    return {'messages': [HumanMessage(content=result)]}


def websearch_node(state: AgentState):

    next_plan = state.get('next_plan', '')
    websearch_prompt = ChatPromptTemplate.from_messages([
            ('system', f'你是一个网络搜索Agent，你需要根据用户的问题，从中提取出用于在网络搜索引擎上搜索的关键字。只需要回答关键字即可，不需要进行解释'),
            ('user', f'当前的用户问题是：'),
            MessagesPlaceholder(variable_name="next_plan"),
        ])
    websearch_chain = (websearch_prompt | llm)
    llm_resp = websearch_chain.invoke({'next_plan': [('user', next_plan)]})
    search_result = do_search(llm_resp.content)
    result_resp = llm.invoke(f'请根据搜索结果进行回答我的问题。\n我的问题是:{next_plan}\n\n搜索结果为:{search_result}')
    return {'messages': [HumanMessage(content=result_resp.content)]}


def composer_node(state: AgentState):
    composer_prompt = ChatPromptTemplate.from_messages([
        ('system', f'你负责根据用户的要求以诗歌的形式进行创作'),
        ('user', f'当前用户的要求是:'),
        MessagesPlaceholder(variable_name="request"),
    ])
    composer_chain = composer_prompt | llm
    llm_resp = composer_chain.invoke({'request': [('user', state.get('next_plan', ''))]})
    return {'messages': [HumanMessage(content=llm_resp.content)]}


def end_node(state: AgentState):
    print('任务完成')


workflow = StateGraph(AgentState)

workflow.add_node('planner', planner)
workflow.add_node('决策节点', decision_node)
workflow.add_node('router', router_node)
workflow.add_node('天气查询模块', weather_node)
workflow.add_node('网络查询模块', websearch_node)
workflow.add_node('诗歌创作模块', composer_node)
workflow.add_node('结束节点', end_node)

workflow.add_edge('planner', '决策节点')
workflow.add_edge('决策节点', 'router')
workflow.add_conditional_edges(
    'router',
    lambda x: x['next'],
    {
        '天气查询模块': '天气查询模块',
        '网络查询模块': '网络查询模块',
        '诗歌创作模块': '诗歌创作模块',
        '结束节点': '结束节点'
    }
)

workflow.add_edge('天气查询模块', '决策节点')
workflow.add_edge('网络查询模块', '决策节点')
workflow.add_edge('诗歌创作模块', '决策节点')
workflow.add_edge('结束节点', END)

workflow.set_entry_point('planner')

graph = workflow.compile()

if '__main__' == __name__:
    print(graph.get_graph().draw_ascii())

    user_query = '查询2024年中国GDP最高的城市，以这个城市今天的天气，创作一首诗'
    enter = {'user_query': user_query, "members": MEMBERS}
    for s in graph.stream(enter):
        if "__end__" not in s and 'router' not in s:
            print(s)
            print('-'*80)

    import datetime
    print(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))




