"""LangGraph agent that answers GAIA benchmark questions (Hugging Face Agents Course, Unit 4).

The agent is a LangGraph ReAct-style graph: an LLM served through the Hugging Face
Inference API, bound to a toolset (web/Wikipedia/arXiv search, arithmetic, GAIA
task-file download, and CSV/Excel analysis). Structure and prompt are adapted from
https://huggingface.co/spaces/fisherman611/gaia-agent, trimmed to tools that need no
extra paid API keys and no arbitrary code execution. See
https://huggingface.co/learn/agents-course/unit4/hands-on for the assignment.
"""

import os
import re
import tempfile
import uuid
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse

import pandas as pd
import requests
from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# Any HF Inference API model that supports tool calling works here. Override via the
# HF_AGENT_MODEL env var without touching code.
HF_MODEL_REPO_ID = os.getenv("HF_AGENT_MODEL", "Qwen/Qwen2.5-Coder-32B-Instruct")

with open(Path(__file__).parent / "system_prompt.txt", "r", encoding="utf-8") as f:
    SYSTEM_PROMPT = f.read()

_FINAL_ANSWER_RE = re.compile(r"final answer\s*:\s*", re.IGNORECASE)


### =============== SEARCH TOOLS =============== ###


@tool
def web_search(query: str) -> str:
    """Search the web (DuckDuckGo) for a query and return a few results.

    Args:
        query: The search query.
    """
    return DuckDuckGoSearchRun().invoke(query)


@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return up to 2 results.

    Args:
        query: The search query.
    """
    docs = WikipediaLoader(query=query, load_max_docs=2).load()
    return "\n\n---\n\n".join(
        f'<Document source="{d.metadata.get("source", "")}"/>\n{d.page_content}\n</Document>'
        for d in docs
    )


@tool
def arxiv_search(query: str) -> str:
    """Search arXiv for a query and return up to 2 results (abstracts truncated).

    Args:
        query: The search query.
    """
    docs = ArxivLoader(query=query, load_max_docs=2).load()
    return "\n\n---\n\n".join(
        f'<Document source="{d.metadata.get("source", "")}"/>\n{d.page_content[:1000]}\n</Document>'
        for d in docs
    )


### =============== MATH TOOLS =============== ###


@tool
def add(a: float, b: float) -> float:
    """Add two numbers.

    Args:
        a: the first number
        b: the second number
    """
    return a + b


@tool
def subtract(a: float, b: float) -> float:
    """Subtract two numbers.

    Args:
        a: the first number
        b: the second number
    """
    return a - b


@tool
def multiply(a: float, b: float) -> float:
    """Multiply two numbers.

    Args:
        a: the first number
        b: the second number
    """
    return a * b


@tool
def divide(a: float, b: float) -> float:
    """Divide two numbers.

    Args:
        a: the numerator
        b: the denominator
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """Get the remainder of a divided by b.

    Args:
        a: the first number
        b: the second number
    """
    return a % b


@tool
def power(a: float, b: float) -> float:
    """Raise a to the power of b.

    Args:
        a: the base
        b: the exponent
    """
    return a**b


@tool
def square_root(a: float) -> float:
    """Get the square root of a non-negative number.

    Args:
        a: the number to get the square root of
    """
    if a < 0:
        raise ValueError("Cannot take the square root of a negative number.")
    return a**0.5


### =============== FILE TOOLS =============== ###


@tool
def download_task_file(task_id: str) -> str:
    """Download the file attached to a GAIA task (if any) and return its text content.

    Only useful when the question references an attached file. Pass the task's task_id.
    Returns decoded text (truncated to 4000 characters) for text-like files, or a short
    description (content type and size) for files that can't be decoded as text.
    """
    try:
        response = requests.get(f"{DEFAULT_API_URL}/files/{task_id}", timeout=30)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        return f"Error downloading file for task {task_id}: {e}"

    try:
        return response.content.decode("utf-8")[:4000]
    except UnicodeDecodeError:
        content_type = response.headers.get("content-type", "unknown")
        return (
            f"File for task {task_id} is binary (content-type: {content_type}, "
            f"{len(response.content)} bytes) and cannot be read as text."
        )


@tool
def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
    """Download a file from a URL to a temporary path, for later analysis.

    Args:
        url: the URL of the file to download.
        filename: optional filename to save as; a random one is used if omitted.
    """
    try:
        if not filename:
            filename = os.path.basename(urlparse(url).path) or f"downloaded_{uuid.uuid4().hex[:8]}"
        filepath = os.path.join(tempfile.gettempdir(), filename)
        response = requests.get(url, stream=True, timeout=30)
        response.raise_for_status()
        with open(filepath, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return f"File downloaded to {filepath}."
    except Exception as e:
        return f"Error downloading file: {e}"


@tool
def analyze_csv_file(file_path: str) -> str:
    """Load a CSV file and return its shape, columns, and summary statistics.

    Args:
        file_path: path to the CSV file (e.g. from download_file_from_url).
    """
    try:
        df = pd.read_csv(file_path)
        return (
            f"{len(df)} rows, {len(df.columns)} columns.\n"
            f"Columns: {', '.join(df.columns)}\n\n"
            f"Summary statistics:\n{df.describe(include='all')}"
        )
    except Exception as e:
        return f"Error analyzing CSV file: {e}"


@tool
def analyze_excel_file(file_path: str) -> str:
    """Load an Excel file and return its shape, columns, and summary statistics.

    Args:
        file_path: path to the .xlsx/.xls file (e.g. from download_file_from_url).
    """
    try:
        df = pd.read_excel(file_path)
        return (
            f"{len(df)} rows, {len(df.columns)} columns.\n"
            f"Columns: {', '.join(df.columns)}\n\n"
            f"Summary statistics:\n{df.describe(include='all')}"
        )
    except Exception as e:
        return f"Error analyzing Excel file: {e}"


def _build_tools():
    return [
        web_search,
        wiki_search,
        arxiv_search,
        add,
        subtract,
        multiply,
        divide,
        modulus,
        power,
        square_root,
        download_task_file,
        download_file_from_url,
        analyze_csv_file,
        analyze_excel_file,
    ]


def _build_llm():
    endpoint = HuggingFaceEndpoint(
        repo_id=HF_MODEL_REPO_ID,
        huggingfacehub_api_token=os.getenv("HF_TOKEN"),
        temperature=0,
        max_new_tokens=1024,
    )
    return ChatHuggingFace(llm=endpoint)


def build_graph():
    """Build the compiled LangGraph agent graph."""
    tools = _build_tools()
    llm_with_tools = _build_llm().bind_tools(tools)

    def assistant(state: MessagesState):
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

    builder = StateGraph(MessagesState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")

    return builder.compile()


def _extract_final_answer(text: str) -> str:
    match = _FINAL_ANSWER_RE.search(text)
    answer = text[match.end():] if match else text
    answer = answer.strip()
    if len(answer) >= 2 and answer[0] == answer[-1] and answer[0] in "\"'":
        answer = answer[1:-1].strip()
    return answer


class GaiaAgent:
    """A LangGraph ReAct agent (HF Inference API LLM + tools) for GAIA questions."""

    def __init__(self):
        self._graph = build_graph()
        print("GaiaAgent initialized.")

    def __call__(self, question: str, task_id: Optional[str] = None) -> str:
        print(f"Agent received question (first 80 chars): {question[:80]}...")

        user_content = question
        if task_id:
            user_content += (
                f"\n\n(task_id: {task_id} - use download_task_file if a file is attached)"
            )

        try:
            result = self._graph.invoke(
                {
                    "messages": [
                        SystemMessage(content=SYSTEM_PROMPT),
                        HumanMessage(content=user_content),
                    ]
                }
            )
            answer = _extract_final_answer(result["messages"][-1].content)
        except Exception as e:
            print(f"Agent error: {e}")
            answer = f"AGENT ERROR: {e}"

        print(f"Agent returning answer (first 80 chars): {answer[:80]}...")
        return answer