# app.py  –  async + progress, keine Cache-Logik
# ------------------------------------------------
import os, asyncio, concurrent.futures, functools
import gradio as gr, requests, pandas as pd
from langchain_core.messages import HumanMessage
from agent import agent_executor     # dein LangGraph-Agent


DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MAX_PAR_TASKS   = 5                  # wie viele Fragen parallel laufen


# ------------------------------------------------------------------
# Sync-Wrapper um den Agent
# ------------------------------------------------------------------
def run_agent_sync(task_id: str, question: str) -> str:
    llm_input = {
        "messages": [HumanMessage(content=question)],
        "task_id":  task_id,
    }
    try:
        result = agent_executor.invoke(llm_input)
        return result["messages"][-1].content.strip()
    except Exception as e:
        return f"AGENT ERROR: {e}"


async def run_agent_async(executor, task_id: str, question: str) -> str:
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(
        executor, functools.partial(run_agent_sync, task_id, question)
    )


# ------------------------------------------------------------------
# Haupt-Callback (async)  –  holt Fragen, verarbeitet parallel
# ------------------------------------------------------------------
async def run_and_submit_all(profile: gr.OAuthProfile | None, progress=gr.Progress()):
    if not profile:
        return "Please login with your HF account.", None
    username = profile.username

    # 1) GAIA-Fragen holen
    q_url = f"{DEFAULT_API_URL}/questions"
    try:
        q_data = requests.get(q_url, timeout=15).json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    progress(0, desc=f"Fetched {len(q_data)} questions – processing …")

    # 2) Parallel ausführen
    answers, log_rows = [], []
    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_PAR_TASKS) as ex:
        tasks = [
            run_agent_async(ex, itm["task_id"], itm["question"])
            for itm in q_data
        ]
        for i, coro in enumerate(asyncio.as_completed(tasks), 1):
            answer = await coro
            task_id = q_data[i-1]["task_id"]
            question = q_data[i-1]["question"]

            answers.append({"task_id": task_id, "submitted_answer": answer})
            log_rows.append({"Task ID": task_id, "Question": question, "Answer": answer})

            progress(i / len(q_data), desc=f"{i}/{len(q_data)} done")

    # 3) Antworten submitten
    submit_url = f"{DEFAULT_API_URL}/submit"
    payload = {
        "username": username,
        "agent_code": f"https://huggingface.co/spaces/{os.getenv('SPACE_ID')}/tree/main",
        "answers": answers,
    }
    try:
        res = requests.post(submit_url, json=payload, timeout=60).json()
        status = (
            f"Submission OK – Score: {res.get('score','?')} % "
            f"({res.get('correct_count','?')}/{res.get('total_attempted','?')})"
        )
    except Exception as e:
        status = f"Submission failed: {e}"

    return status, pd.DataFrame(log_rows)


# ------------------------------------------------------------------
# Gradio-UI
# ------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Fast GAIA Agent Runner  (Async + Progress)")
    gr.LoginButton()
    run_btn = gr.Button("Run & Submit")

    out_status = gr.Textbox(label="Status / Score", lines=3, interactive=False)
    out_table  = gr.DataFrame(label="Answers", wrap=True)

    run_btn.click(run_and_submit_all, outputs=[out_status, out_table])

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