from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
import os
from tools import search_tool
from langgraph.graph import StateGraph, START, END, add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from typing import TypedDict, Annotated
from langchain_core.messages import AnyMessage, SystemMessage

SYSTEM_PROMPT = (
                """You are an Answer Summarizer AI assistant.
                Here is the question: {state.get(state["messages"][0].content)}
                Format the answer with the following rules:
                - YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
                - If the question asked for a number, write only the number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. Write numbers using digits.
                - If the question asked for a string, don't use articles, no abbreviations (e.g: write "Saint" rather than "St."), and write the digits in plain text unless specified otherwise.
                - If the question 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.
                WRITE ONLY THE ANSWER, no sentence, no explanation, no context"""
)

def create_agent():
    # Generate the chat interface, including the tools
    from langchain_mistralai import ChatMistralAI
    
    chat = ChatMistralAI(
        model="mistral-large-latest",
        temperature=0,
        max_retries=2
    )
    tools = [search_tool]
    chat_with_tools = chat.bind_tools(tools)

    # Generate the AgentState and Agent graph
    class AgentState(TypedDict):
        messages: Annotated[list[AnyMessage], add_messages]

    def assistant(state: AgentState):
        messages_for_model = [
            SystemMessage(content=SYSTEM_PROMPT),
            *state["messages"],
        ]
        return {
            "messages": [chat_with_tools.invoke(messages_for_model)],
        }

    ## The graph
    builder = StateGraph(AgentState)

    # Define nodes: these do the work
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))

    # Define edges: these determine how the control flow moves
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges(
        "assistant",
        # If the latest message requires a tool, route to tools
        # Otherwise, provide a direct response
        tools_condition,
    )
    builder.add_edge("tools", "assistant")
    agent = builder.compile()
    return agent