"""
agent.py — LangChain ReAct agent for the GAIA benchmark.

Tools included:
  • DuckDuckGo web search (free, no API key)
  • Wikipedia lookup
  • Python REPL (math, data manipulation)
  • File download + reading (text, CSV, PDF via pdfminer)
  • Image understanding via HuggingFace Inference API (free tier)
  • ArXiv search
"""

import os
import io
import re
import tempfile
import traceback
from typing import Optional

import requests

# LangChain core
from langchain.agents import AgentExecutor, create_react_agent
from langchain.tools import Tool, tool
from langchain_core.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEndpoint

# Community tools
from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_experimental.tools.python.tool import PythonREPLTool


# ---------------------------------------------------------------------------
# 1. LLM — free HuggingFace Inference Endpoint
# ---------------------------------------------------------------------------

def get_llm():
    """
    Use a capable open model via HF Inference API (free tier).
    Qwen2.5-72B-Instruct is a strong publicly available model.
    You can swap for meta-llama/Meta-Llama-3-70B-Instruct, etc.
    Requires HF_TOKEN env var (free account works).
    """
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise EnvironmentError(
            "HF_TOKEN environment variable not set. "
            "Add it in your HuggingFace Space secrets."
        )

    llm = HuggingFaceEndpoint(
        repo_id="Qwen/Qwen2.5-72B-Instruct",
        huggingfacehub_api_token=hf_token,
        task="text-generation",
        max_new_tokens=1024,
        temperature=0.1,
        do_sample=False,
        repetition_penalty=1.1,
    )
    return llm


# ---------------------------------------------------------------------------
# 2. Tool definitions
# ---------------------------------------------------------------------------

# -- Web search --
def make_search_tool():
    search = DuckDuckGoSearchRun()
    return Tool(
        name="web_search",
        func=search.run,
        description=(
            "Search the web for current information using DuckDuckGo. "
            "Use this for facts, recent events, people, places, or anything "
            "that requires up-to-date knowledge. Input: a search query string."
        ),
    )


# -- Wikipedia --
def make_wikipedia_tool():
    wiki = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper(top_k_results=3, doc_content_chars_max=3000))
    return Tool(
        name="wikipedia",
        func=wiki.run,
        description=(
            "Look up encyclopedic information on Wikipedia. "
            "Best for well-known topics, historical facts, science, biographies. "
            "Input: a topic or question string."
        ),
    )


# -- Python REPL --
def make_python_tool():
    repl = PythonREPLTool()
    return Tool(
        name="python_repl",
        func=repl.run,
        description=(
            "Execute Python code for calculations, data processing, string manipulation, "
            "logic, and analysis. pandas, math, re, json, csv, datetime are available. "
            "Input: valid Python code as a string. Always print() the result you need."
        ),
    )


# -- File reader --
@tool
def read_file_from_url(url: str) -> str:
    """
    Download a file from a URL and return its text content.
    Supports: plain text (.txt, .py, .json, .csv, .md), PDF, and basic image description.
    Input: a URL string pointing to the file.
    """
    try:
        resp = requests.get(url, timeout=30)
        resp.raise_for_status()
        content_type = resp.headers.get("content-type", "")

        # PDF
        if "pdf" in content_type or url.lower().endswith(".pdf"):
            try:
                from pdfminer.high_level import extract_text
                with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
                    f.write(resp.content)
                    tmp_path = f.name
                text = extract_text(tmp_path)
                os.unlink(tmp_path)
                return text[:5000] if text else "Could not extract PDF text."
            except ImportError:
                return "pdfminer not available. Install pdfminer.six to read PDFs."

        # CSV
        if "csv" in content_type or url.lower().endswith(".csv"):
            import csv
            decoded = resp.content.decode("utf-8", errors="replace")
            lines = decoded.splitlines()
            return "\n".join(lines[:50])  # first 50 rows

        # Excel
        if url.lower().endswith((".xlsx", ".xls")):
            try:
                import pandas as pd
                from io import BytesIO
                df = pd.read_excel(BytesIO(resp.content))
                return df.to_string(max_rows=50)
            except Exception as e:
                return f"Could not read Excel file: {e}"

        # Image — describe via HF
        if "image" in content_type or url.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".gif")):
            return describe_image_bytes(resp.content)

        # Default: text
        return resp.content.decode("utf-8", errors="replace")[:5000]

    except Exception as e:
        return f"Error reading file from {url}: {e}"


# -- Image understanding via HF Inference API --
def describe_image_bytes(image_bytes: bytes) -> str:
    """Use HF Inference API to caption an image."""
    hf_token = os.getenv("HF_TOKEN", "")
    headers = {}
    if hf_token:
        headers["Authorization"] = f"Bearer {hf_token}"

    # Use Salesforce BLIP image captioning (free on HF Inference API)
    api_url = "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-large"
    try:
        response = requests.post(api_url, headers=headers, data=image_bytes, timeout=30)
        result = response.json()
        if isinstance(result, list) and result:
            return result[0].get("generated_text", "No caption generated.")
        return str(result)
    except Exception as e:
        return f"Image description failed: {e}"


@tool
def describe_image_from_url(url: str) -> str:
    """
    Download an image from a URL and return a text description of its contents.
    Use this when a question refers to an image file.
    Input: a direct URL to an image (jpg, png, webp, etc.).
    """
    try:
        resp = requests.get(url, timeout=30)
        resp.raise_for_status()
        return describe_image_bytes(resp.content)
    except Exception as e:
        return f"Could not describe image: {e}"


# -- ArXiv search --
@tool
def arxiv_search(query: str) -> str:
    """
    Search ArXiv for scientific papers. Use for questions about research, ML papers,
    physics, mathematics, computer science publications.
    Input: a search query string.
    Returns: titles, authors, and abstracts of top results.
    """
    try:
        import urllib.parse
        encoded = urllib.parse.quote(query)
        url = f"https://export.arxiv.org/api/query?search_query=all:{encoded}&start=0&max_results=3"
        resp = requests.get(url, timeout=15)
        resp.raise_for_status()

        # Parse simple XML
        text = resp.text
        entries = re.findall(r"<entry>(.*?)</entry>", text, re.DOTALL)
        results = []
        for entry in entries:
            title = re.search(r"<title>(.*?)</title>", entry, re.DOTALL)
            summary = re.search(r"<summary>(.*?)</summary>", entry, re.DOTALL)
            authors = re.findall(r"<name>(.*?)</name>", entry)
            t = title.group(1).strip() if title else "?"
            s = summary.group(1).strip()[:500] if summary else ""
            a = ", ".join(authors[:3])
            results.append(f"Title: {t}\nAuthors: {a}\nAbstract: {s}")
        return "\n\n---\n\n".join(results) if results else "No results found."
    except Exception as e:
        return f"ArXiv search failed: {e}"


# -- Final answer formatter --
@tool
def final_answer(answer: str) -> str:
    """
    Use this tool to submit the final answer to the question.
    Input: the exact final answer string.
    """
    return answer


# ---------------------------------------------------------------------------
# 3. ReAct Prompt
# ---------------------------------------------------------------------------

REACT_PROMPT = PromptTemplate.from_template(
    """You are an expert AI assistant solving questions from the GAIA benchmark.
GAIA tests real-world question answering that requires reasoning, web search, file reading, and multi-step problem solving.

You have access to the following tools:
{tools}

Use this format strictly:

Question: the input question you must answer
Thought: reason step-by-step about what to do
Action: the action to take, must be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (repeat Thought/Action/Action Input/Observation as needed)
Thought: I now know the final answer
Final Answer: the exact answer to the question

Rules:
- Be precise. GAIA expects exact answers (numbers, names, dates, etc.).
- Use web_search and wikipedia for factual lookups.
- Use python_repl for any calculations, unit conversions, or data analysis.
- Use read_file_from_url if a file URL is provided.
- Use describe_image_from_url if an image URL is provided.
- Use arxiv_search for scientific paper questions.
- If the question asks for a number, return just the number.
- If the question asks for a name, return just the name.
- Do not add explanation to Final Answer — just the answer.
- Limit reasoning to what is necessary.

Begin!

Question: {input}
Thought:{agent_scratchpad}"""
)


# ---------------------------------------------------------------------------
# 4. Build agent
# ---------------------------------------------------------------------------

def build_agent():
    llm = get_llm()

    tools = [
        make_search_tool(),
        make_wikipedia_tool(),
        make_python_tool(),
        read_file_from_url,
        describe_image_from_url,
        arxiv_search,
    ]

    agent = create_react_agent(
        llm=llm,
        tools=tools,
        prompt=REACT_PROMPT,
    )

    executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        max_iterations=10,
        max_execution_time=120,
        handle_parsing_errors=True,
        return_intermediate_steps=False,
    )

    def run(question: str) -> str:
        try:
            result = executor.invoke({"input": question})
            return str(result.get("output", "No answer produced.")).strip()
        except Exception as e:
            print(f"Agent error: {traceback.format_exc()}")
            return f"Error: {e}"

    return run