import os
import re
import requests
import pandas as pd
import gradio as gr
from smolagents import CodeAgent, InferenceClientModel, DuckDuckGoSearchTool, VisitWebpageTool

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

INSTRUCTIONS = """You are answering GAIA benchmark questions. Use tools when needed.
Reply with ONLY the final answer, no explanation and no 'FINAL ANSWER' prefix.
- Numbers: plain digits, no commas, no units unless asked.
- Strings: no articles or abbreviations unless asked.
- Lists: comma-separated, in the order requested.
Answers are graded by exact match."""


class BasicAgent:
    def __init__(self):
        model = InferenceClientModel(
            model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
            token=os.getenv("HF_TOKEN"),
        )
        self.agent = CodeAgent(
            tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
            model=model,
            additional_authorized_imports=["pandas", "numpy", "requests", "re", "json", "datetime", "openpyxl"],
            instructions=INSTRUCTIONS,
            max_steps=10,
        )

    def __call__(self, question, task_id=None, file_name=None):
        prompt = question
        if file_name:
            try:
                r = requests.get(f"{API}/files/{task_id}", timeout=30)
                path = f"/tmp/{file_name}"
                with open(path, "wb") as f:
                    f.write(r.content)
                prompt += f"\n\nA file is attached at: {path}"
            except Exception:
                pass
        try:
            ans = str(self.agent.run(prompt)).strip()
            ans = re.sub(r"^(final answer:?\s*)", "", ans, flags=re.I).strip()
            return ans.rstrip(".")
        except Exception:
            return "unknown"


def run_and_submit_all(profile: gr.OAuthProfile | None):
    if not profile:
        return "Please log in to Hugging Face first.", None
    username = profile.username
    space_id = os.getenv("SPACE_ID")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

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

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

    answers, log = [], []
    for item in questions:
        tid = item["task_id"]
        try:
            ans = agent(item["question"], tid, item.get("file_name"))
        except Exception as e:
            ans = "unknown"
        answers.append({"task_id": tid, "submitted_answer": ans})
        log.append({"Task ID": tid, "Question": item["question"], "Answer": ans})

    payload = {"username": username.strip(), "agent_code": agent_code, "answers": answers}
    try:
        res = requests.post(f"{API}/submit", json=payload, timeout=120).json()
        status = (
            f"Submission successful!\nUser: {res.get('username')}\n"
            f"Score: {res.get('score')}% "
            f"({res.get('correct_count')}/{res.get('total_attempted')} correct)\n"
            f"{res.get('message', '')}"
        )
    except Exception as e:
        status = f"Submission failed: {e}"
    return status, pd.DataFrame(log)


with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent - Final Assignment")
    gr.LoginButton()
    btn = gr.Button("Run Evaluation & Submit All Answers")
    status = gr.Textbox(label="Status", lines=6, interactive=False)
    table = gr.DataFrame(label="Questions and Answers", wrap=True)
    btn.click(run_and_submit_all, outputs=[status, table])

if __name__ == "__main__":
    demo.launch()