import os
import re
import tempfile

import gradio as gr
import pandas as pd
import requests
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    tool,
)

try:
    from smolagents import InferenceClientModel as HFModel
except ImportError:  # older smolagents versions
    from smolagents import HfApiModel as HFModel

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct")


# --- Tools ---
@tool
def get_task_file(task_id: str) -> str:
    """Download the file attached to a task (if any) and return its content or local path.

    Args:
        task_id: The task_id of the question that has an attached file.
    """
    try:
        r = requests.get(f"{DEFAULT_API_URL}/files/{task_id}", timeout=60)
        r.raise_for_status()
    except Exception as e:
        return f"No file could be downloaded for this task: {e}"

    disposition = r.headers.get("content-disposition", "")
    match = re.search(r'filename="?([^";]+)"?', disposition)
    name = match.group(1) if match else f"{task_id}.bin"
    path = os.path.join(tempfile.gettempdir(), name)
    with open(path, "wb") as f:
        f.write(r.content)

    lower = name.lower()
    try:
        if lower.endswith((".xlsx", ".xls")):
            sheets = pd.read_excel(path, sheet_name=None)
            return "\n\n".join(
                f"Sheet '{s}':\n{df.to_string()}" for s, df in sheets.items()
            )[:20000]
        if lower.endswith((".txt", ".csv", ".py", ".json", ".md")):
            return r.content.decode("utf-8", errors="replace")[:20000]
    except Exception as e:
        return f"File saved at {path} but could not be parsed: {e}"
    return (
        f"File saved at {path} (type: {name.split('.')[-1]}). "
        "It is a binary file (image/audio/other); handle it with Python code."
    )


SYSTEM_HINT = """You are solving GAIA-style questions. Rules:
- Use tools (web search, visiting pages, reading attached files, Python) to find the answer.
- Your final answer must be ONLY the answer: a number, a short string, or a comma-separated list.
- Do NOT write 'FINAL ANSWER', explanations, units unless asked, or full sentences.
- Numbers: no thousands separators, no currency/percent symbols unless asked.
- Strings: no articles (a, the) and no abbreviations unless asked.
- Lists: comma-separated, in the order requested.
- If the question mentions an attached file, call get_task_file with the given task_id.
"""


# --- Agent ---
class BasicAgent:
    def __init__(self):
        model = HFModel(model_id=MODEL_ID, token=os.getenv("HF_TOKEN"))
        self.agent = CodeAgent(
            tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), get_task_file],
            model=model,
            additional_authorized_imports=[
                "pandas", "numpy", "requests", "re", "json",
                "datetime", "math", "bs4", "collections", "itertools",
            ],
            max_steps=12,
        )

    def __call__(self, question: str, task_id: str = "", has_file: bool = False) -> str:
        prompt = f"{SYSTEM_HINT}\n\ntask_id: {task_id}\n"
        if has_file:
            prompt += "This task has an attached file (use get_task_file).\n"
        prompt += f"\nQuestion: {question}"
        answer = str(self.agent.run(prompt)).strip()
        answer = re.sub(r"(?i)^\s*final answer\s*:\s*", "", answer)
        return answer.strip().rstrip(".") if len(answer) > 1 else answer


# --- Run & submit ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")
    if not profile:
        return "Please log in to Hugging Face with the button.", None
    username = profile.username

    api_url = DEFAULT_API_URL
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None

    try:
        response = requests.get(f"{api_url}/questions", timeout=30)
        response.raise_for_status()
        questions = response.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    results_log, answers_payload = [], []
    for item in questions:
        task_id = item.get("task_id")
        question = item.get("question")
        has_file = bool(item.get("file_name"))
        if not task_id or question is None:
            continue
        try:
            answer = agent(question, task_id=task_id, has_file=has_file)
        except Exception as e:
            answer = f"AGENT ERROR: {e}"
        answers_payload.append({"task_id": task_id, "submitted_answer": answer})
        results_log.append(
            {"Task ID": task_id, "Question": question, "Submitted Answer": answer}
        )

    if not answers_payload:
        return "Agent produced no answers.", pd.DataFrame(results_log)

    submission = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }
    try:
        resp = requests.post(f"{api_url}/submit", json=submission, timeout=120)
        resp.raise_for_status()
        data = resp.json()
        status = (
            f"Submission successful!\n"
            f"User: {data.get('username')}\n"
            f"Score: {data.get('score')}% "
            f"({data.get('correct_count')}/{data.get('total_attempted')} correct)\n"
            f"Message: {data.get('message')}"
        )
        return status, pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission failed: {e}", pd.DataFrame(results_log)


# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("# Agents Course - Final Assignment")
    gr.Markdown(
        "1. Log in with your Hugging Face account.\n"
        "2. Click the button to run the agent on all questions and submit.\n"
        "This can take several minutes."
    )
    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])

if __name__ == "__main__":
    demo.launch(debug=True, share=False)