import os
import gradio as gr
import pandas as pd
import requests
from transformers import pipeline

# --- Local LLM Agent ---
class LocalLLMAgent:
    def __init__(self):
        print("Loading local model...")
        self.generator = pipeline(
            "text-generation",
            model="distilgpt2",  # small, fast
            device=-1  # CPU only
        )

    def __call__(self, question: str) -> str:
        try:
            outputs = self.generator(
                question,
                max_length=50,
                do_sample=False,
                num_return_sequences=1,
            )
            return outputs[0]["generated_text"].strip()
        except Exception as e:
            print(f"[LocalLLMAgent Error] {e}")
            return "Error: Could not generate answer."


# --- Dummy Fallback Agent ---
class DummyAgent:
    def __call__(self, question: str) -> str:
        return "Error: Could not generate answer."


# --- GAIA Submission Logic ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    GAIA_API = "https://agents-course-unit4-scoring.hf.space"
    space_id = os.getenv("SPACE_ID", "XeroDN/final_assignment")
    username = profile.username if profile else "anonymous"

    try:
        agent = LocalLLMAgent()
    except Exception as e:
        print("Falling back to DummyAgent:", e)
        agent = DummyAgent()

    # Fetch GAIA questions
    try:
        qres = requests.get(f"{GAIA_API}/questions", timeout=15)
        qres.raise_for_status()
        questions = qres.json()
    except Exception as e:
        return f"❌ Failed to fetch questions: {e}", pd.DataFrame()

    # Generate answers
    results_log = []
    answers_payload = []
    for item in questions:
        q = item.get("question")
        tid = item.get("task_id")
        if not q or not tid:
            continue
        answer = agent(q)
        results_log.append({"Task ID": tid, "Question": q, "Submitted Answer": answer})
        answers_payload.append({"task_id": tid, "submitted_answer": answer})

    if not answers_payload:
        return "⚠️ No answers generated.", pd.DataFrame()

    # Submit answers
    try:
        submission = {
            "username": username,
            "agent_code": f"https://huggingface.co/spaces/{space_id}/tree/main",
            "answers": answers_payload
        }
        sres = requests.post(f"{GAIA_API}/submit", json=submission, timeout=60)
        sres.raise_for_status()
        result = sres.json()
        score = result.get("score", "?")
        correct = result.get("correct_count", "?")
        total = result.get("total_attempted", "?")
        summary = f"✅ Submission Successful!\nUser: {username}\nScore: {score}% ({correct}/{total})"
        return summary, pd.DataFrame(results_log)
    except Exception as e:
        return f"❌ Submission failed: {e}", pd.DataFrame(results_log)

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 GAIA Benchmark Agent Submission (Local LLM, no API key)")
    gr.Markdown("Click below to log in and run your agent on the GAIA benchmark.")

    login_btn = gr.LoginButton()
    run_btn = gr.Button("🚀 Run Evaluation & Submit All Answers")
    status = gr.Textbox(label="Status", lines=6)
    table = gr.DataFrame(label="Results")

    run_btn.click(fn=run_and_submit_all, inputs=[login_btn], outputs=[status, table])

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