import functools, operator

from datetime import date

from typing import Annotated, Any, Dict, List, Optional, Sequence, Tuple, TypedDict, Union

from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next: str

def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="messages"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    
    agent = create_openai_tools_agent(llm, tools, prompt)
    executor = AgentExecutor(agent=agent, tools=tools)
    
    return executor

def agent_node(state, agent, name):
    result = agent.invoke(state)
    return {"messages": [HumanMessage(content=result["output"], name=name)]}

@tool
def today_tool(text: str) -> str:
    """Returns today's date. Use this for any questions related to knowing today's date. 
       The input should always be an empty string, and this function will always return today's date. 
       Any date mathematics should occur outside this function."""
    return (str(date.today()) + "\n\nIf you have completed all tasks, respond with FINAL ANSWER.")
    
def create_graph(model, max_tokens, temperature, topic):
    tavily_tool = TavilySearchResults(max_results=10)
    
    members = ["Content Planner", "Content Writer"]
    options = ["FINISH"] + members
   
    system_prompt = (
        "You are a Manager tasked with managing a conversation between the "
        "following agent(s): {members}. Given the following user request, "
        "respond with the agent to act next. Each agent will perform a "
        "task and respond with their results and status. When finished, "
        "respond with FINISH."
    )

    function_def = {
        "name": "route",
        "description": "Select the next role.",
        "parameters": {
            "title": "routeSchema",
            "type": "object",
            "properties": {
                "next": {
                    "title": "Next",
                    "anyOf": [
                        {"enum": options},
                    ],
                }
            },
            "required": ["next"],
        },
    }
    
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="messages"),
            (
                "system",
                "Given the conversation above, who should act next? "
                "Or should we FINISH? Select one of: {options}.",
            ),
        ]
    ).partial(options=str(options), members=", ".join(members))
    
    llm = ChatOpenAI(model=model, max_tokens=max_tokens, temperature=temperature)
    
    supervisor_chain = (
        prompt
        | llm.bind_functions(functions=[function_def], function_call="route")
        | JsonOutputFunctionsParser()
    )

    content_planner_agent = create_agent(llm, [tavily_tool], system_prompt=
                                    "You are a Content Planner working on planning a blog article "
              "about the topic: " + topic + "."
              "You collect information that helps the "
              "audience learn something "
              "and make informed decisions. "
              "Your work is the basis for "
              "the Content Writer to write an article on this topic.")
    content_planner_node = functools.partial(agent_node, agent=content_planner_agent, name="Content Planner")

    content_writer_agent = create_agent(llm, [today_tool], system_prompt=
                                    "You are a Content Writer working on writing "
              "a new opinion piece about the topic: " + topic + ". "
              "You base your writing on the work of "
              "the Content Planner, who provides an outline "
              "and relevant context about the topic. "
              "You follow the main objectives and "
              "direction of the outline, "
              "as provide by the Content Planner. "
              "You also provide objective and impartial insights "
              "and back them up with information "
              "provide by the Content Planner. "
              "You acknowledge in your opinion piece "
              "when your statements are opinions "
              "as opposed to objective statements.")
    content_writer_node = functools.partial(agent_node, agent=content_writer_agent, name="Content Writer")
    
    workflow = StateGraph(AgentState)

    workflow.add_node("Manager", supervisor_chain)
    workflow.add_node("Content Planner", content_planner_node)
    workflow.add_node("Content Writer", content_writer_node)

    for member in members:
        workflow.add_edge(member, "Manager")

    conditional_map = {k: k for k in members}
    conditional_map["FINISH"] = END
    workflow.add_conditional_edges("Manager", lambda x: x["next"], conditional_map)
    
    workflow.set_entry_point("Manager")
    
    return workflow.compile()

def run_multi_agent(llm, max_tokens, temperature, topic):
    graph = create_graph(llm, max_tokens, temperature, topic)
    
    result = graph.invoke({
        "messages": [
            HumanMessage(content=topic)
        ]
    })
    
    article = result['messages'][-1].content
    
    print("===")
    print(article)
    print("===")
    
    return article