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

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    InferenceClientModel
)

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


# -----------------------------
# Smart Agent
# -----------------------------
class BasicAgent:
    def __init__(self):

        print("Initializing Smart Agent...")

        # Web Search Tool
        search_tool = DuckDuckGoSearchTool()

        # Free Hugging Face Model
        model = InferenceClientModel(
            model_id="meta-llama/Llama-3.1-8B-Instruct"
        )

        # Main Agent
        self.agent = CodeAgent(
            tools=[search_tool],
            model=model,
            add_base_tools=True,
            max_steps=5
        )

    def __call__(self, question: str) -> str:

        print(f"Question: {question}")

        prompt = f"""
You are a GAIA benchmark assistant.

IMPORTANT RULES:
- Return ONLY the final answer
- Do NOT explain your reasoning
- Do NOT write 'FINAL ANSWER'
- Keep answers short and exact
- If the answer is a number, return only the number
- If the answer is text, return only the text

Question:
{question}
"""

        try:
            response = self.agent.run(prompt)

            answer = str(response).strip()

            print(f"Agent answer: {answer}")

            return answer

        except Exception as e:
            print(f"Error while solving question: {e}")

            return "Error"


# -----------------------------
# Main Evaluation Function
# -----------------------------
def run_and_submit_all(profile: gr.OAuthProfile | None):

    space_id = os.getenv("SPACE_ID")

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

    api_url = DEFAULT_API_URL

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

    # -----------------------------
    # Create Agent
    # -----------------------------
    try:
        agent = BasicAgent()

    except Exception as e:
        return f"Error initializing agent: {e}", None

    # -----------------------------
    # Space Code URL
    # -----------------------------
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    print(agent_code)

    # -----------------------------
    # Fetch Questions
    # -----------------------------
    try:
        response = requests.get(
            questions_url,
            timeout=30
        )

        response.raise_for_status()

        questions_data = response.json()

        print(f"Fetched {len(questions_data)} questions")

    except Exception as e:
        return f"Error fetching questions: {e}", None

    # -----------------------------
    # Run Agent
    # -----------------------------
    results_log = []

    answers_payload = []

    for item in questions_data:

        task_id = item.get("task_id")

        question_text = item.get("question")

        if not task_id or question_text is None:
            continue

        try:
            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
            })

        except Exception as e:

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

    # -----------------------------
    # Submit Answers
    # -----------------------------
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }

    try:

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

        response.raise_for_status()

        result_data = response.json()

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

        results_df = pd.DataFrame(results_log)

        return final_status, results_df

    except Exception as e:

        results_df = pd.DataFrame(results_log)

        return f"Submission Failed: {e}", results_df


# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks() as demo:

    gr.Markdown("# GAIA Agent Evaluation")

    gr.Markdown(
        """
Login with Hugging Face and run your AI agent on GAIA questions.
"""
    )

    gr.LoginButton()

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

    status_output = gr.Textbox(
        label="Status",
        lines=8
    )

    results_table = gr.DataFrame(
        label="Agent Results"
    )

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


# -----------------------------
# Launch App
# -----------------------------
if __name__ == "__main__":

    print("Starting GAIA Agent App...")

    demo.launch(
        debug=True,
        share=False
    )