import os
import re

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

import agent_tools
import spaces


@spaces.GPU
def _gpu_stub():
    # ZeroGPU Spaces need at least one @spaces.GPU function to start; this agent does not use a GPU.
    return None

API_URL = agent_tools.API_URL
MODEL_CANDIDATES = [
    "Qwen/Qwen3-235B-A22B-Instruct-2507",
    "deepseek-ai/DeepSeek-V3",
    "Qwen/Qwen2.5-72B-Instruct",
    "meta-llama/Llama-3.3-70B-Instruct",
]

INSTRUCTIONS = """You answer GAIA benchmark questions. Rules:
- Your final_answer must be ONLY the answer itself: no explanation, no 'FINAL ANSWER' prefix, no trailing period.
- Numbers: plain digits, no commas or units unless the question asks for them. Lists: comma-separated in the order requested.
- If the question prompt names an attached file, call get_task_file(task_id, file_name) first, then use the matching reader tool.
- Search the web and open pages to verify facts; prefer Wikipedia (wikipedia_page) and primary sources. Never guess when a lookup is possible.
- If a question is written backwards or in code, decode it first, then answer it.
"""


def pick_model(token: str):
    key = os.getenv("LLM_API_KEY") or os.getenv("GEMINI_API_KEY")
    if key:
        from smolagents import OpenAIServerModel

        base = os.getenv("LLM_BASE_URL") or "https://generativelanguage.googleapis.com/v1beta/openai/"
        names = [m for m in [os.getenv("LLM_MODEL"), os.getenv("GEMINI_MODEL")] if m] or ["gemini-2.5-flash"]
        last = None
        for mid in names:
            try:
                model = OpenAIServerModel(model_id=mid, api_base=base, api_key=key)
                model([{"role": "user", "content": [{"type": "text", "text": "ping"}]}])
                return model
            except Exception as e:  # noqa: BLE001
                last = e
        raise RuntimeError(f"LLM unavailable ({base}): {last}")
    last = None
    for mid in MODEL_CANDIDATES:
        try:
            model = InferenceClientModel(model_id=mid, token=token)
            model([{"role": "user", "content": [{"type": "text", "text": "ping"}]}])
            return model
        except Exception as e:  # noqa: BLE001
            last = e
    raise RuntimeError(f"No inference model available: {last}")


def build_agent(token: str) -> CodeAgent:
    agent_tools.HF_TOKEN["value"] = token
    return CodeAgent(
        tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), *agent_tools.ALL_TOOLS],
        model=pick_model(token),
        additional_authorized_imports=["pandas", "numpy", "re", "json", "math", "datetime", "itertools", "collections", "requests", "bs4"],
        max_steps=10,
        instructions=INSTRUCTIONS,
    )


def run_with_retry(agent, prompt, tries=4):
    import time

    for n in range(tries):
        try:
            return agent.run(prompt, reset=True)
        except Exception as e:  # noqa: BLE001
            if n == tries - 1 or not any(k in str(e) for k in ("429", "503", "rate", "quota", "overloaded")):
                raise
            time.sleep(30 * (n + 1))


def clean(ans) -> str:
    s = str(ans).strip()
    s = re.sub(r"(?i)^\s*final answer\s*[:\-]\s*", "", s)
    return s.strip().rstrip(".") if len(s.split()) < 8 else s


def run_and_submit_all(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None, progress=gr.Progress()):
    if not profile or not oauth_token:
        return "Please log in with the Hugging Face button first.", None
    username = profile.username
    space_id = os.getenv("SPACE_ID", f"{username}/Final_Assignment_Template")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    try:
        questions = requests.get(f"{API_URL}/questions", timeout=30).json()
        agent = build_agent(oauth_token.token)
    except Exception as e:  # noqa: BLE001
        return f"Setup failed: {e}", None

    answers, log = [], []
    for i, item in enumerate(questions):
        progress((i, len(questions)), desc=f"Question {i + 1}/{len(questions)}")
        tid, q, fname = item["task_id"], item["question"], item.get("file_name") or ""
        prompt = q
        if fname:
            prompt += f"\n\n[Attached file: task_id={tid}, file_name={fname}]"
        try:
            ans = clean(run_with_retry(agent, prompt))
        except Exception as e:  # noqa: BLE001
            ans = ""
            print(f"task {tid} failed: {e}")
        answers.append({"task_id": tid, "submitted_answer": ans})
        log.append({"Task ID": tid, "Question": q[:120], "Submitted Answer": ans})

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


with gr.Blocks() as demo:
    gr.Markdown("# GAIA Final Assignment Agent\nLog in, then run. The agent answers all 20 questions and submits them.")
    gr.LoginButton()
    btn = gr.Button("Run Evaluation & Submit All Answers")
    out = gr.Textbox(label="Status", lines=5, interactive=False)
    table = gr.DataFrame(label="Questions and answers", wrap=True)
    btn.click(run_and_submit_all, outputs=[out, table])

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