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

from smolagents import LiteLLMModel, CodeAgent, DuckDuckGoSearchTool
from gaia_tools import ReverseTextTool, RunPythonFileTool, download_server


# ==============================
# System Prompt
# ==============================

SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question.
Report your thoughts, and finish your answer with just the answer — no prefixes.

Your answer should be:
• A number
OR
• Few words
OR
• Comma-separated list

Rules:

If number:
- No commas
- No units

If string:
- No articles
- No abbreviations
- Write digits as words

Tool Rules:

1. Use only provided tools.
2. Use one tool at a time.
3. If reversed question → use ReverseTextTool.
4. If .py file → use RunPythonFileTool.
5. For downloads → use download_server.
"""


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


# ==============================
# Agent Class (NVIDIA Version)
# ==============================

class MyAgent:

    def __init__(self):

        nvidia_api_key = os.getenv("NVIDIA_API_KEY")

        if not nvidia_api_key:
            raise ValueError(
                "NVIDIA_API_KEY not set in environment variables."
            )

        # NVIDIA MiniMax model via LiteLLM
        self.model = LiteLLMModel(

            model_id="openai/minimaxai/minimax-m2.5",

            api_key=nvidia_api_key,

            api_base="https://integrate.api.nvidia.com/v1",

            system_prompt=SYSTEM_PROMPT

        )

        self.agent = CodeAgent(

            tools=[
                DuckDuckGoSearchTool(),
                ReverseTextTool,
                RunPythonFileTool,
                download_server
            ],

            model=self.model,

            add_base_tools=True,

        )

    def __call__(self, question: str) -> str:
        return self.agent.run(question)


# ==============================
# Main Evaluation Function
# ==============================

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.", None


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


    try:
        agent = MyAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None


    # Fetch questions

    try:

        response = requests.get(
            questions_url,
            timeout=15
        )

        response.raise_for_status()

        questions_data = response.json()

    except Exception as e:

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


    results_log = []
    answers_payload = []


    # Run agent

    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"AGENT ERROR: {e}"
            })


    if not answers_payload:
        return "Agent did not return answers.", pd.DataFrame(results_log)


    submission_data = {

        "username": profile.username.strip(),

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

        "answers": answers_payload

    }


    # Submit answers

    try:

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

        response.raise_for_status()

        result_data = response.json()

        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/"
            f"{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )

        return final_status, 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("# NVIDIA MiniMax Agent Runner 🚀")

    gr.Markdown("""
**Instructions**

1. Add NVIDIA API key in Secrets  
2. Login to HuggingFace  
3. Click Run  
""")

    gr.LoginButton()

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

    status_output = gr.Textbox(
        label="Submission Result",
        lines=5,
        interactive=False
    )

    results_table = gr.DataFrame(
        label="Results",
        wrap=True
    )

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


if __name__ == "__main__":

    print("🔧 App starting...")

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