import os
import gradio as gr
import requests
import pandas as pd
from openai import OpenAI

# =============================
# CONSTANTS (DO NOT CHANGE)
# =============================
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

SYSTEM_PROMPT = """
You are a general AI assistant.
Answer the question and finish your response with:

FINAL ANSWER: <answer>

Rules:
- If number: no commas, no units unless specified
- If string: no articles, no abbreviations
- If list: comma-separated, minimal words
"""

# =============================
# MODEL
# =============================
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def llm_answer(question: str) -> str:
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": question}
        ],
        temperature=0
    )

    text = response.choices[0].message.content

    # Extract FINAL ANSWER only
    if "FINAL ANSWER:" in text:
        return text.split("FINAL ANSWER:")[-1].strip()

    return text.strip()

# =============================
# GAIA PIPELINE
# =============================
def run_and_submit_all(profile: gr.OAuthProfile | None):
    if not profile:
        return "❌ Please login to Hugging Face", None

    username = profile.username.strip()
    space_id = os.getenv("SPACE_ID")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    # Fetch questions
    questions = requests.get(f"{DEFAULT_API_URL}/questions").json()

    answers = []
    logs = []

    for q in questions:
        try:
            ans = llm_answer(q["question"])
        except Exception as e:
            ans = "I don't know"

        answers.append({
            "task_id": q["task_id"],
            "submitted_answer": ans
        })

        logs.append({
            "Task ID": q["task_id"],
            "Question": q["question"],
            "Answer": ans
        })

    payload = {
        "username": username,
        "agent_code": agent_code,
        "answers": answers
    }

    res = requests.post(f"{DEFAULT_API_URL}/submit", json=payload).json()

    status = (
        f"✅ Submission Successful!\n"
        f"User: {res.get('username')}\n"
        f"Score: {res.get('score')}%\n"
        f"Correct: {res.get('correct_count')}/{res.get('total_attempted')}\n"
        f"Message: {res.get('message')}"
    )

    return status, pd.DataFrame(logs)

# =============================
# UI
# =============================
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 GAIA Level-1 Agent (Unit-4)")
    gr.LoginButton()

    run_btn = gr.Button("Run Evaluation & Submit")
    status = gr.Textbox(label="Submission Result", lines=5)
    table = gr.Dataframe(label="Questions & Answers")

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

demo.launch()