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
from langchain_mistralai import ChatMistralAI

# SYSTEM_PROMPT = (
#     "You are a general AI assistant. "
#     "I will ask you a question. "
#     "Only write your answer, nothing more, following the instructions below. "
#     "Your 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."
# )

# SYSTEM_PROMPT = """
# Answer the question with the exact final answer only.

# Never output explanations, reasoning, labels, markdown, quotes, or extra punctuation.

# Allowed outputs only:
# - a single number
# - a single short string
# - a comma-separated list of numbers and/or strings

# Mandatory formatting:
# - If list: use commas with NO spaces: a,b,c
# - If number: no thousands separators, no units, no symbols unless explicitly requested
# - If string: as few words as possible, no articles, no abbreviations unless explicitly requested

# Final verification before output:
# - remove any extra text
# - remove any spaces around commas
# - remove any unrequested units or symbols

# Output only the answer.
# """

SYSTEM_PROMPT = """
You are answering a GAIA evaluation item.

Your task is to produce the final answer in the exact required format.

Output rules:
- Output the answer only.
- Do not add any explanation, reasoning, prefix, suffix, label, or punctuation other than what belongs to the answer itself.
- Do not use markdown.
- Do not use quotes.
- Do not add a trailing period.
- Do not add a newline before or after the answer.
- If the answer is a single number, output only that number.
- If the answer is a single string, output only that string.
- If the answer is a comma-separated list, output the items separated by commas with NO spaces before or after commas.

Formatting constraints:
- Numbers:
  - Do not use thousands separators.
  - Do not use commas inside numbers.
  - Do not add units unless explicitly requested.
  - Do not add symbols such as $, %, or ° unless explicitly requested.
- Strings:
  - Use as few words as possible.
  - Do not use articles.
  - Do not use abbreviations unless explicitly requested.
  - For city names, country names, person names, etc., write the full form.
  - Write digits in words only if explicitly required by the question or instructions.
- Comma-separated lists:
  - Use this exact format: item1,item2,item3
  - NO spaces anywhere around commas.
  - Apply the number/string rules to each item individually.

Priority order:
1. Follow the user’s question.
2. Follow these formatting rules exactly.
3. When in doubt, prefer the shortest valid answer.

Before sending the answer, perform this silent check:
- Is there any extra text? If yes, remove it.
- Are there any spaces around commas? If yes, remove them.
- Are there any articles in a string answer? If yes, remove them.
- Are there any abbreviations in a string answer? If yes, expand them unless explicitly allowed.
- Are there any units or symbols not explicitly requested? If yes, remove them.

Return only the final answer.
"""

def create_agent():
    # # Generate the chat interface, including the tools
    # llm = HuggingFaceEndpoint(
    #     # repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    #     # repo_id="deepseek-ai/DeepSeek-R1",
    #     repo_id="Qwen/Qwen3-Coder-480B-A35B-Instruct",
    #     # repo_id="meta-llama/Llama-3.1-8B-Instruct",
    #     # provider="auto",
    #     # provider="together",
    #     # repo_id="moonshotai/Kimi-K2.5",
    #     # provider="auto",
    #     huggingfacehub_api_token=os.getenv("HF_TOKEN"),
    #     server_kwargs={
    #         "bill_to": "HornetSecurity"
    #     },
    # )

    chat = ChatMistralAI(
            model="mistral-large-2512",
            api_key=os.environ["MISTRAL_API_KEY"],
            temperature=1,
    )
    # chat = ChatHuggingFace(llm=llm, verbose=True)
    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