"""
app.py — HuggingFace Agents Course Unit 4 Final Assignment
Uses Google Gemini (gemini-2.0-flash) as the agent model.

Deploy to HF Spaces:
  - SDK: Gradio
  - Add secret: Gemini

Get your free Gemini API key at: https://aistudio.google.com/apikey
Scoring API docs: https://agents-course-unit4-scoring.hf.space/docs
Course leaderboard: https://huggingface.co/spaces/agents-course/Students_leaderboard
"""

import os
import re
import time

import gradio as gr
import requests
from google import genai
from google.genai import types

# ── Constants ─────────────────────────────────────────────────────────────────

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MODEL = "gemini-2.0-flash"

SYSTEM_PROMPT = (
    "You are a precise AI assistant. Answer the question as concisely as possible. "
    "Your final answer must be a number, as few words as possible, or a comma-separated "
    "list of numbers and/or strings. "
    "Do not include units (e.g. $, %) unless explicitly asked. "
    "Do not use articles or abbreviations for place names. "
    "Write numbers as digits, not words, unless asked otherwise. "
    "Reply with ONLY the answer — no explanation, no preamble."
)

# ── Agent ─────────────────────────────────────────────────────────────────────

def run_agent(question: str, task_id: str, api_url: str) -> str:
    """Call Gemini to answer one GAIA question. Attaches file if one exists."""
    api_key = os.environ.get("Gemini", "")
    if not api_key:
        return "[ERROR] Gemini secret is not set in this Space."

    client = genai.Client(api_key=api_key)

    # Build message parts — try to fetch any associated file
    parts = []
    try:
        file_resp = requests.get(f"{api_url}/files/{task_id}", timeout=10)
        if file_resp.status_code == 200 and file_resp.content:
            import base64
            content_type = file_resp.headers.get("content-type", "")

            if "pdf" in content_type or "image" in content_type:
                media_type = content_type.split(";")[0].strip()
                parts.append(types.Part(
                    inline_data=types.Blob(
                        mime_type=media_type,
                        data=base64.standard_b64encode(file_resp.content).decode(),
                    )
                ))
            else:
                try:
                    text_content = file_resp.content.decode("utf-8", errors="replace")
                    parts.append(types.Part(text=f"[Attached file content]\n{text_content}\n"))
                except Exception:
                    pass
    except Exception:
        pass  # No file or fetch failed — continue without it

    parts.append(types.Part(text=question))

    config = types.GenerateContentConfig(
        system_instruction=SYSTEM_PROMPT,
        tools=[types.Tool(google_search=types.GoogleSearch())],
        safety_settings=[
            types.SafetySetting(category="HARM_CATEGORY_HARASSMENT",        threshold="BLOCK_NONE"),
            types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH",       threshold="BLOCK_NONE"),
            types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
            types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
        ],
    )

    for attempt in range(5):
        try:
            response = client.models.generate_content(
                model=MODEL,
                contents=parts,
                config=config,
            )
            answer = response.text.strip()
            answer = re.sub(r"(?i)^final answer:\s*", "", answer).strip()
            return answer
        except Exception as e:
            msg = str(e)
            if "429" in msg or "RESOURCE_EXHAUSTED" in msg:
                wait = 30 * (attempt + 1)  # 30s, 60s, 90s, 120s, 150s
                time.sleep(wait)
            elif attempt == 4:
                return f"[ERROR] {e}"
            else:
                time.sleep(5)

    return "[ERROR] All retries failed."


# ── Gradio logic ──────────────────────────────────────────────────────────────

def run_and_submit(profile: gr.OAuthProfile | None, *args):
    """Fetch questions → run agent → submit answers → return results."""

    if profile is None:
        return "Please log in with your Hugging Face account first.", None

    username = profile.username
    space_id = os.environ.get("SPACE_ID", "")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else "N/A"
    api_url = DEFAULT_API_URL

    # Fetch questions
    status_lines = [f"Logged in as: **{username}**", "Fetching questions…"]
    try:
        resp = requests.get(f"{api_url}/questions", timeout=15)
        resp.raise_for_status()
        questions = resp.json()
    except Exception as e:
        return f"Failed to fetch questions: {e}", None

    if not questions:
        return "No questions returned by the API.", None

    status_lines.append(f"Fetched {len(questions)} questions. Running agent…")

    # Run agent on each question
    results_log = []
    answers_payload = []

    for i, item in enumerate(questions, 1):
        task_id = item.get("task_id", "")
        question = item.get("question", "")

        answer = run_agent(question, task_id, api_url)
        time.sleep(4)  # stay under free-tier RPM limit (15 req/min)

        answers_payload.append({"task_id": task_id, "submitted_answer": answer})
        results_log.append({
            "Task ID": task_id,
            "Question": question[:80] + ("…" if len(question) > 80 else ""),
            "Answer": answer,
        })
        status_lines.append(f"[{i}/{len(questions)}] {task_id} → {answer[:60]}")

    # Submit answers
    status_lines.append("Submitting answers…")
    submit_payload = {
        "username": username,
        "agent_code": agent_code,
        "answers": answers_payload,
    }

    try:
        submit_resp = requests.post(f"{api_url}/submit", json=submit_payload, timeout=60)
        submit_resp.raise_for_status()
        result = submit_resp.json()

        score = result.get("score", "N/A")
        correct = result.get("correct_count", "?")
        total = result.get("total_attempted", len(questions))
        message = result.get("message", "")

        status_lines.append(f"\n**Score: {score}% ({correct}/{total} correct)**\n{message}")
    except Exception as e:
        status_lines.append(f"Submission failed: {e}")

    return "\n\n".join(status_lines), results_log


# ── Gradio UI ─────────────────────────────────────────────────────────────────

with gr.Blocks(title="GAIA Agent — Unit 4 Final Assignment") as demo:
    gr.Markdown(
        """
        # GAIA Agent — HF Agents Course Unit 4
        This agent fetches GAIA benchmark questions, answers them using **Gemini 2.0 Flash**
        with web search, and submits your score to the
        [course leaderboard](https://huggingface.co/spaces/agents-course/Students_leaderboard).

        **Before running:** make sure `Gemini` is set as a secret in your Space settings.
        Get a free key at [aistudio.google.com/apikey](https://aistudio.google.com/apikey).
        """
    )

    login_btn = gr.LoginButton()

    run_btn = gr.Button("Run Evaluation & Submit", variant="primary")

    status_box = gr.Markdown(label="Status")

    results_table = gr.Dataframe(
        headers=["Task ID", "Question", "Answer"],
        datatype=["str", "str", "str"],
        label="Results",
        wrap=True,
    )

    run_btn.click(
        fn=run_and_submit,
        inputs=[login_btn],
        outputs=[status_box, results_table],
    )

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