from typing import TypedDict, Annotated, Sequence
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage, BaseMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph, MessagesState, END
from langgraph.prebuilt import tools_condition
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_core.runnables import RunnableConfig

from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun, ArxivQueryRun
from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
from langchain_community.tools.wikidata.tool import WikidataAPIWrapper, WikidataQueryRun

from langchain_openai import ChatOpenAI

from tools import GetYouTubeTranscriptTool, ImageRecognitionTool
from dotenv import load_dotenv
import os

load_dotenv()
HUGGINGFACEHUB_API_TOKEN = os.getenv("HF_TOKEN")

# Define your state class if needed
class AgentState(TypedDict):
    """The state of the agent."""

    messages: Annotated[Sequence[BaseMessage], add_messages]



def build_graph():
    """
    Build and return the compiled LangGraph Runnable agent.
    """

    def call_model(
        state: AgentState,
        config: RunnableConfig,
    ):
        system_prompt = SystemMessage("You are a general AI assistant. I will ask you a question. Report only your final answer without the thoughts or any other text. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. Also be direct when doing a tool call. Try using other tools before using the DuckDuckGoTool. If you can't find informations, answer based on your personal knowledge.")
        
        response = model.invoke([system_prompt] + state["messages"], config)
        # We return a list, because this will get added to the existing list
        return {"messages": [response]}

    def should_continue(state: AgentState):
        messages = state["messages"]
        last_message = messages[-1]
        print(last_message)
        # If there is no function call, then we finish
        if not last_message.tool_calls:
            return "end"
        # Otherwise if there is, we continue
        else:
            return "continue"
        
    model = ChatOpenAI(model="o1")
    image_model = ChatOpenAI(model="gpt-4o")

    tools = [GetYouTubeTranscriptTool(), 
             WikidataQueryRun(api_wrapper=WikidataAPIWrapper()), 
             DuckDuckGoSearchRun(),
             ArxivQueryRun(api_wrapper=ArxivAPIWrapper()),
             WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()), 
             ImageRecognitionTool(hf_endpoint=image_model)]
    model = model.bind_tools(tools)

    tool_node = ToolNode(tools)
    
    workflow = StateGraph(AgentState)

    # Define the two nodes we will cycle between
    workflow.add_node("agent", call_model)
    workflow.add_node("tools", tool_node)
    
    # Set the entrypoint as `agent`
    # This means that this node is the first one called
    workflow.set_entry_point("agent")
    
    # We now add a conditional edge
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {
            # If `tools`, then we call the tool node.
            "continue": "tools",
            "end": END
        },
    )
    
    workflow.add_edge("tools", "agent")
    
    graph = workflow.compile()

    return graph