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

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


class BasicAgent:
    def __init__(self):
        hf_token = os.getenv("HF_TOKEN")
        if not hf_token:
            raise ValueError("HF_TOKEN missing")

        self.headers = {
            "Authorization": f"Bearer {hf_token}"
        }

        self.api_url = "https://router.huggingface.co/v1/chat/completions"
        self.model = "meta-llama/Llama-3.1-8B-Instruct"

        print("Agent Ready")

    def __call__(self, question: str):

        prompt = f"""
You are a GAIA solving agent.

Rules:
- Think step by step internally
- Give ONLY final answer
- No explanation
- No extra words

Question:
{question}

Final Answer:
"""

        payload = {
            "model": self.model,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "temperature": 0
        }

        try:
            r = requests.post(
                self.api_url,
                headers=self.headers,
                json=payload,
                timeout=120
            )

            r.raise_for_status()

            result = r.json()

            return result["choices"][0]["message"]["content"].strip()

        except Exception as e:
            print("Error:", e)
            return "ERROR"


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

    space_id = os.getenv("SPACE_ID")

    if not profile:
        return "Login required", None

    username = profile.username

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

    agent = BasicAgent()

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

    try:
        r = requests.get(questions_url, timeout=30)
        r.raise_for_status()
        questions = r.json()
    except Exception as e:
        return f"Fetch error: {e}", None

    answers = []
    logs = []

    for item in questions:
        task_id = item.get("task_id")
        question = item.get("question")

        if not task_id or not question:
            continue

        ans = agent(question)

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

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

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

    try:
        r = requests.post(submit_url, json=payload, timeout=120)
        r.raise_for_status()
        data = r.json()

        return (
            f"Score: {data.get('score')}%\n"
            f"Correct: {data.get('correct_count')}/{data.get('total_attempted')}\n"
            f"{data.get('message')}",
            pd.DataFrame(logs)
        )

    except Exception as e:
        return f"Submit error: {e}", pd.DataFrame(logs)


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

    gr.LoginButton()

    btn = gr.Button("Run & Submit")

    out = gr.Textbox(label="Result", lines=6)
    table = gr.DataFrame(label="Logs")

    btn.click(run_and_submit_all, outputs=[out, table])


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