"""GAIA Lite evaluation runner with answer cache, attachment download, and split submit."""
import json
import os
import tempfile
import time

import gradio as gr
import pandas as pd
import requests

from agent import run_agent, build_graph
from files_util import DEFAULT_API_URL, download_scoring_file

CACHE_PATH = os.getenv("GAIA_LITE_ANSWERS_CACHE", os.path.join(os.path.dirname(__file__), "answers_cache.json"))
MAX_RETRIES = 3


def load_cache() -> dict:
    if not os.path.exists(CACHE_PATH):
        return {}
    try:
        with open(CACHE_PATH, "r", encoding="utf-8") as handle:
            return json.load(handle)
    except json.JSONDecodeError:
        return {}


def save_cache(cache: dict) -> None:
    with open(CACHE_PATH, "w", encoding="utf-8") as handle:
        json.dump(cache, handle, ensure_ascii=False, indent=2)


class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
        self.graph = build_graph(use_retriever=False)

    def __call__(self, question: str, file_paths=None) -> str:
        print(f"Agent received question (first 80 chars): {question[:80]}...")
        result = run_agent(question, file_paths=file_paths, graph=self.graph)
        print(f"Extracted FINAL ANSWER: {result['final']}")
        return result["final"]


def fetch_questions(api_url: str) -> list:
    response = requests.get(f"{api_url}/questions", timeout=30)
    response.raise_for_status()
    data = response.json()
    if not data:
        raise ValueError("Fetched questions list is empty.")
    return data


def run_one_question(agent: BasicAgent, item: dict, work_dir: str, api_url: str) -> str:
    question_text = item.get("question") or ""
    task_id = item["task_id"]
    file_name = item.get("file_name") or ""
    file_paths = []
    if file_name:
        saved = download_scoring_file(task_id, file_name, dest_dir=work_dir, api_url=api_url)
        if saved:
            file_paths.append(saved)
        else:
            question_text += (
                f"\n\nNote: an attachment named {file_name} was expected but could not be downloaded."
            )
    return agent(question_text, file_paths=file_paths)


def invoke_with_retry(agent: BasicAgent, item: dict, work_dir: str, api_url: str) -> str:
    delay = 5
    last_error = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            return run_one_question(agent, item, work_dir, api_url)
        except Exception as exc:
            last_error = exc
            print(f"Attempt {attempt}/{MAX_RETRIES} failed for {item.get('task_id')}: {exc}")
            if attempt < MAX_RETRIES:
                time.sleep(delay)
                delay = min(delay * 2, 60)
    return f"AGENT ERROR: {last_error}"


def results_dataframe(questions_data: list, cache: dict) -> pd.DataFrame:
    rows = []
    for item in questions_data:
        task_id = item.get("task_id")
        cached = cache.get(task_id, {})
        rows.append(
            {
                "Task ID": task_id,
                "Question": item.get("question"),
                "File": item.get("file_name") or "",
                "Submitted Answer": cached.get("submitted_answer", ""),
            }
        )
    return pd.DataFrame(rows)


def run_evaluation(profile: gr.OAuthProfile | None):
    if not profile:
        return "Please log in to Hugging Face first.", None

    api_url = DEFAULT_API_URL
    try:
        questions_data = fetch_questions(api_url)
    except Exception as exc:
        return f"Error fetching questions: {exc}", None

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

    cache = load_cache()
    work_dir = tempfile.mkdtemp(prefix="gaia_lite_files_")
    print(f"Running agent on {len(questions_data)} questions; cache={CACHE_PATH}")

    for item in questions_data:
        task_id = item.get("task_id")
        if not task_id or item.get("question") is None:
            continue
        if task_id in cache and not str(cache[task_id].get("submitted_answer", "")).startswith("AGENT ERROR"):
            print(f"Skip cached {task_id}")
            continue
        answer = invoke_with_retry(agent, item, work_dir, api_url)
        cache[task_id] = {
            "question": item.get("question"),
            "file_name": item.get("file_name") or "",
            "submitted_answer": answer,
        }
        save_cache(cache)

    df = results_dataframe(questions_data, cache)
    done = sum(1 for item in questions_data if item.get("task_id") in cache)
    return f"Evaluation finished. Cached {done}/{len(questions_data)} answers at {CACHE_PATH}", df


def submit_cached(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"
    api_url = DEFAULT_API_URL

    try:
        questions_data = fetch_questions(api_url)
    except Exception as exc:
        return f"Error fetching questions: {exc}", None

    cache = load_cache()
    answers_payload = []
    for item in questions_data:
        task_id = item.get("task_id")
        cached = cache.get(task_id)
        if not cached:
            continue
        answer = cached.get("submitted_answer", "")
        if str(answer).startswith("AGENT ERROR"):
            continue
        answers_payload.append({"task_id": task_id, "submitted_answer": answer})

    if not answers_payload:
        return "No cached answers to submit. Run evaluation first.", results_dataframe(questions_data, cache)

    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }
    try:
        response = requests.post(f"{api_url}/submit", json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        status = (
            f"Submission successful.\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Submitted {len(answers_payload)} answers.\n"
            f"Message: {result_data.get('message', '')}"
        )
        return status, results_dataframe(questions_data, cache)
    except requests.HTTPError as exc:
        detail = exc.response.text[:500] if exc.response is not None else str(exc)
        return f"Submission failed: {detail}", results_dataframe(questions_data, cache)
    except Exception as exc:
        return f"Submission failed: {exc}", results_dataframe(questions_data, cache)


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # GAIA Lite Evaluation

        1. Log in with Hugging Face.
        2. **Run evaluation** answers every question (skips successful cache entries, writes `answers_cache.json`).
        3. **Submit cached answers** posts the cache to the scoring API.

        Attachments are downloaded from `/files/{task_id}` when `file_name` is present.
        """
    )
    gr.LoginButton()
    with gr.Row():
        run_button = gr.Button("Run evaluation (no submit)", variant="primary")
        submit_button = gr.Button("Submit cached answers")
    status_output = gr.Textbox(label="Status", lines=6, interactive=False)
    results_table = gr.DataFrame(label="Questions and answers", wrap=True)
    run_button.click(fn=run_evaluation, outputs=[status_output, results_table])
    submit_button.click(fn=submit_cached, outputs=[status_output, results_table])


if __name__ == "__main__":
    print("Launching GAIA Lite evaluation UI...")
    demo.launch(debug=True, share=False)