"""
HF Agents Course — Unit 4
GAIA Level 1 Agent (LangGraph + Qwen3.5-35B-A3B)
"""

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage
from typing import TypedDict, Annotated
import operator

# =========================
# 基础配置
# =========================
load_dotenv()

AGICTO_API_KEY = os.getenv("AGICTO_API_KEY")
AGICTO_BASE_URL = os.getenv("AGICTO_BASE_URL", "https://agicto.com/v1")

llm = ChatOpenAI(
    model="qwen3.5-35b-a3b",
    openai_api_key=AGICTO_API_KEY,
    openai_api_base=AGICTO_BASE_URL,
    temperature=0.1,
    max_tokens=1024,
)

# =========================
# System Prompt（GAIA 专用）
# =========================
SYSTEM_PROMPT = (
    "You are a GAIA benchmark agent.\n"
    "Think step-by-step internally.\n"
    "Do NOT explain your reasoning.\n"
    "Output ONLY the final exact answer.\n"
    "No units, no punctuation, no extra words.\n"
    "If the question includes a file URL, use its content."
)

# =========================
# LangGraph State
# =========================
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]

def build_graph():
    def run_agent(state: AgentState) -> AgentState:
        response = llm.invoke(state["messages"])
        return {"messages": [response]}

    graph = StateGraph(AgentState)
    graph.add_node("agent", run_agent)
    graph.set_entry_point("agent")
    graph.add_edge("agent", END)
    return graph.compile()

graph = build_graph()

# =========================
# ✅ HF 官方接口
# =========================
def agent(question: str, files: list[str] | None = None) -> str:
    """
    必须保留此函数签名，HF Unit 4 评分器会直接调用。
    """
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]

    result = graph.invoke({"messages": messages})
    answer = result["messages"][-1].content

    # ✅ GAIA 硬性规则
    return answer.strip().splitlines()[0]