"""GAIA Agent using Free HuggingFace Models - FINAL FIXED VERSION"""
import logging
from typing import Dict
from langchain_community.llms import HuggingFacePipeline
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import START, StateGraph, MessagesState
from transformers import pipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("agent")

def build_graph():
    """Builds a simple agent graph using a Q&A HuggingFace model."""

    # Use a well-supported, general Q&A model (NO auth needed)
    def get_hf_pipeline():
        hf_pipe = pipeline(
            "text2text-generation",
            model="google/flan-t5-base",
            device=-1,  # CPU
            max_length=128
        )
        return HuggingFacePipeline(pipeline=hf_pipe)

    llm = get_hf_pipeline()

    def agent(state: MessagesState) -> Dict:
        question = state["messages"][-1].content
        logger.info(f"Processing: {question[:70]}...")

        # Handle basic direct math: Only if clean digits/operators
        import re
        if re.match(r'^\s*\d+\s*[\+\-\*/]\s*\d+\s*$', question):
            try:
                answer = str(eval(question, {"__builtins__": {}}))
                return {"messages": [AIMessage(content=f"FINAL ANSWER: {answer}")]}
            except Exception:
                pass  # fallback to LLM

        # Otherwise, use LLM for everything else
        prompt = f"Answer as concisely as possible: {question}\nAnswer:"
        try:
            response = llm.invoke(prompt)
            # Get answer string
            if hasattr(response, 'content'):
                answer = response.content
            else:
                answer = str(response)
            # Remove 'Answer:' prefix and keep only the first line
            answer = answer.replace("Answer:", "").strip().split("\n")[0]
            # If the answer is empty, fallback
            if not answer:
                answer = "Unknown"
            return {"messages": [AIMessage(content=f"FINAL ANSWER: {answer}")]}
        except Exception as e:
            logger.error(f"Agent error: {e}")
            return {"messages": [AIMessage(content="FINAL ANSWER: Error processing question")]}

    # Graph setup
    builder = StateGraph(MessagesState)
    builder.add_node("agent", agent)
    builder.add_edge(START, "agent")
    return builder.compile()