import os
import gradio as gr
import requests
import pandas as pd

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


# --- FIXED BASIC AGENT ---
from transformers import pipeline

class BasicAgent:
    def __init__(self):
        self.model = pipeline("text-generation", model="gpt2")

    def __call__(self, question: str) -> str:
        prompt = (
            "Answer ONLY with the final short result.\n"
            "No explanation.\n\n"
            f"Question: {question}\nAnswer:"
        )

        output = self.model(prompt, max_new_tokens=50, do_sample=False)[0]["generated_text"]

        # extract only answer part
        return output.split("Answer:")[-1].strip().split("\n")[0]


def run_and_submit_all(profile: gr.OAuthProfile | None):

    space_id = os.getenv("SPACE_ID")

    if profile:
        username = profile.username
        print(f"User logged in: {username}")
    else:
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # 1. Agent
    agent = BasicAgent()

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    # 2. Fetch questions
    response = requests.get(questions_url, timeout=15)
    questions_data = response.json()

    results_log = []
    answers_payload = []

    # 3. Run agent
    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")

        submitted_answer = agent(question_text)

        answers_payload.append({
            "task_id": task_id,
            "submitted_answer": submitted_answer
        })

        results_log.append({
            "Task ID": task_id,
            "Question": question_text,
            "Submitted Answer": submitted_answer
        })

    # 4. Submit
    submission_data = {
        "username": username,
        "agent_code": agent_code,
        "answers": answers_payload
    }

    response = requests.post(submit_url, json=submission_data, timeout=60)
    result_data = response.json()

    final_status = (
        f"Submission Successful!\n"
        f"User: {result_data.get('username')}\n"
        f"Score: {result_data.get('score')}%\n"
        f"Message: {result_data.get('message')}"
    )

    return final_status, pd.DataFrame(results_log)


# --- UI ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")

    gr.LoginButton()

    run_button = gr.Button("Run Evaluation & Submit All Answers")

    status_output = gr.Textbox(label="Status", lines=5)
    results_table = gr.DataFrame(label="Results")

    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table]
    )

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