"""
app.py: Hugging Face Agents Course, Unit 4 final assignment (faster version).

Changes from the template:
- questions run in parallel (WORKERS at a time); the rate limiter in agent.py keeps
  the total number of model calls inside the free tier
- each question has a time cap, so one stuck question can't hold up the whole run
- answers are cached in answers_cache.json: if a run stops halfway, the next run
  only answers what's missing
- a placeholder @spaces.GPU function, needed only when the Space runs on ZeroGPU
"""

import os
import json
import time
import concurrent.futures as cf

import gradio as gr
import pandas as pd
import requests

import agent as agent_module
from agent import BasicAgent, check_model

try:                       # ZeroGPU checks for one @spaces.GPU function at startup
    import spaces

    @spaces.GPU
    def gpu_placeholder():
        return None
except ImportError:        # CPU Spaces don't have (or need) the spaces package
    pass

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
WORKERS = 2                  # questions answered at the same time (the agent's rate limiter paces Groq anyway)
RUN_TIME_CAP = 3600          # seconds for the whole run; unfinished questions are submitted as "unknown"
CACHE_FILE = "answers_cache.json"


def load_cache():
    try:
        with open(CACHE_FILE) as f:
            return json.load(f)
    except (OSError, ValueError):
        return {}


def save_cache(cache):
    with open(CACHE_FILE, "w") as f:
        json.dump(cache, f, indent=2)


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

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

    # 0. check the API key and quota with one tiny call before spending time on 20 questions
    problem = check_model()
    if problem:
        return ("Test call to the model failed, so nothing was run or submitted.\n"
                f"{problem[:800]}\n\n"
                "Usual causes: Groq's daily free limit used up (wait until tomorrow), a wrong or missing "
                "GROQ_API_KEY secret, or a model name that changed (edit MODEL_ID in agent.py)."), None

    # 1. fetch questions
    try:
        r = requests.get(f"{api_url}/questions", timeout=20)
        r.raise_for_status()
        questions = r.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None
    if not questions:
        return "The question list is empty.", None

    # 2. answer them in parallel, reusing cached answers
    cache = load_cache() if use_cache else {}
    todo = [q for q in questions if q.get("task_id") and q["task_id"] not in cache]
    print(f"{len(questions)} questions, {len(questions) - len(todo)} from cache, {len(todo)} to answer")
    done = 0

    start = time.time()
    pool = cf.ThreadPoolExecutor(max_workers=WORKERS)
    futures = {pool.submit(agent, q["question"], q["task_id"]): q for q in todo}
    try:
        for fut in cf.as_completed(futures, timeout=RUN_TIME_CAP):
            q = futures[fut]
            try:
                answer = fut.result()
            except Exception as e:
                print(f"[{q['task_id']}] error: {e}")
                answer = "unknown"
            if answer and answer != "unknown":      # only real answers are cached
                cache[q["task_id"]] = answer
                save_cache(cache)
            done += 1
            print(f"[{done}/{len(todo)}] {time.time() - start:.0f}s  {answer!r}")
    except cf.TimeoutError:
        print("Time cap reached; unfinished questions are submitted as 'unknown'.")
    pool.shutdown(wait=False, cancel_futures=True)

    # 3. build the payload and the results table
    answers_payload, log = [], []
    for q in questions:
        answer = cache.get(q["task_id"], "unknown")
        answers_payload.append({"task_id": q["task_id"], "submitted_answer": answer})
        log.append({"Task ID": q["task_id"], "Question": q["question"], "Submitted Answer": answer,
                    "Error": agent.last_error.get(q["task_id"], "")})
    results_df = pd.DataFrame(log)

    # 4. submit
    submission = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    try:
        r = requests.post(f"{api_url}/submit", json=submission, timeout=60)
        r.raise_for_status()
        d = r.json()
        status = (
            f"Model: {agent_module.MODEL_ID}\n"
            f"Submission successful in {time.time() - start:.0f}s\n"
            f"User: {d.get('username')}\n"
            f"Score: {d.get('score', 'N/A')}% "
            f"({d.get('correct_count', '?')}/{d.get('total_attempted', '?')} correct)\n"
            f"Message: {d.get('message', '')}"
        )
    except requests.exceptions.HTTPError as e:
        status = f"Submission failed: {e.response.status_code} {e.response.text[:500]}"
    except Exception as e:
        status = f"Submission failed: {e}"
    return status, results_df


with gr.Blocks() as demo:
    gr.Markdown("# Agents Course: Unit 4 final assignment")
    gr.Markdown(
        "1. Log in with your Hugging Face account.\n"
        "2. Click **Run Evaluation & Submit All Answers**. Questions run in parallel; "
        "with the cache on, a second run only answers the questions that are still missing."
    )
    gr.LoginButton()
    use_cache = gr.Checkbox(value=True, label="Reuse answers from the previous run")
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run status / submission result", lines=6, interactive=False)
    results_table = gr.DataFrame(label="Questions and agent answers", wrap=True)
    run_button.click(fn=run_and_submit_all, inputs=[use_cache], outputs=[status_output, results_table])

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