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

from smolagents import CodeAgent, DuckDuckGoSearchTool, TransformersModel

DEFAULT_API_URL = os.getenv("DEFAULT_API_URL", "https://agents-course-unit4-scoring.hf.space")

# You can override the model via the MODEL_ID Space variable.
# Use a small instruct/chat model so the agent can plan and call tools reliably.
DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "HuggingFaceTB/SmolLM2-1.7B-Instruct")


# DuckDuckGo-based Wikipedia-style tool (just rename + description)
class WikipediaSearchTool(DuckDuckGoSearchTool):
    name = "wikipedia_search"
    description = "Search Wikipedia-style information using DuckDuckGo backend."


class BasicAgent:
    def __init__(self):
        # Use the official smolagents TransformersModel wrapper so the agent
        # can properly handle chat messages and tool calls, as in the course.
        model = TransformersModel(
            model_id=DEFAULT_MODEL_ID,
            max_new_tokens=512,
        )
        tools = [DuckDuckGoSearchTool(), WikipediaSearchTool()]
        self.agent = CodeAgent(
            tools=tools,
            model=model,
            additional_authorized_imports=["math", "datetime"],
        )

    def __call__(self, question: str) -> str:
        # smolagents versions differ:
        # - some expect agent.run("question")
        # - some accept messages list
        try:
            out = self.agent.run(question)
        except TypeError:
            out = self.agent.run([{"role": "user", "content": question}])
        return str(out).strip()


def _get_space_repo_id() -> str:
    # Preferred: full id
    space_id = (os.getenv("SPACE_ID") or "").strip()
    if space_id and "/" in space_id:
        return space_id

    # Fallback: author + repo name
    author = (os.getenv("SPACE_AUTHOR_NAME") or "").strip()
    repo = (os.getenv("SPACE_REPO_NAME") or "").strip()
    if author and repo:
        return f"{author}/{repo}"

    # Fallback: host sometimes exists like "monkminer-final-assignment-template.hf.space"
    # (Not always reversible reliably, so only last resort)
    return ""


def run_and_submit_all(profile: gr.OAuthProfile | None) -> tuple[str, pd.DataFrame]:
    if not profile:
        return "Please log in with your Hugging Face account.", pd.DataFrame()

    space_id = _get_space_repo_id()
    if not space_id or "/" not in space_id:
        return (
            "Error: Could not detect Space repo id. "
            "Make sure you're running inside a HF Space (and not locally).",
            pd.DataFrame(),
        )

    username = profile.username
    agent_code_url = f"https://huggingface.co/spaces/{space_id}/blob/main/app.py"

    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", pd.DataFrame()

    # Fetch questions
    try:
        resp = requests.get(f"{DEFAULT_API_URL}/questions", timeout=30)
        resp.raise_for_status()
        questions = resp.json()
        if not questions:
            return "No questions received from the API.", pd.DataFrame()
    except Exception as e:
        return f"Error fetching questions: {e}", pd.DataFrame()

    results_log = []
    answers_payload = []

    for item in questions:
        task_id = item.get("task_id")
        question = item.get("question")
        if not task_id or not question:
            continue

        try:
            submitted_answer = agent(question)
            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
            results_log.append(
                {"Task ID": task_id, "Question": question, "Submitted Answer": submitted_answer}
            )
        except Exception as e:
            results_log.append(
                {"Task ID": task_id, "Question": question, "Submitted Answer": f"AGENT ERROR: {e}"}
            )

    if not answers_payload:
        return "Agent failed to generate any answers.", pd.DataFrame(results_log)

    submission_data = {
        "username": username,
        "agent_code": agent_code_url,
        "answers": answers_payload,
    }

    # Submit
    try:
        res = requests.post(f"{DEFAULT_API_URL}/submit", json=submission_data, timeout=120)
        res.raise_for_status()
        data = res.json()
        summary = (
            f" Submission Successful!\n"
            f"User: {data.get('username')}\n"
            f"Score: {data.get('score', 'N/A')}% "
            f"({data.get('correct_count', '?')}/{data.get('total_attempted', '?')})\n"
            f"Message: {data.get('message', 'No message received.')}"
        )
        return summary, pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)


with gr.Blocks() as demo:
    gr.Markdown("# Agent Evaluation Runner")
    gr.Markdown(
        "1) Clone this Space and edit `app.py`.\n"
        "2) Log in with your HF account.\n"
        "3) Click the button to evaluate and submit."
    )

    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Submission Status", lines=6, interactive=False)
    results_table = gr.DataFrame(label="Agent Responses")

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

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