import os
import gradio as gr
import requests
import pandas as pd
from smolagents.core import Agent, tool
from duckduckgo_search import DDGS
from transformers import pipeline

# --- Tool Definitions ---
@tool
class WebSearchTool:
    name = "web_search"
    description = "Search the web for up-to-date factual information."
    def use(self, query: str) -> str:
        with DDGS() as ddgs:
            results = ddgs.text(query)
            output = [f"{r['title']} - {r['href']}" for r in results[:3]]
            return "\n".join(output) if output else "No relevant results found."

@tool
class CiteTool:
    name = "cite"
    description = "Add citation to a given answer with a valid URL."
    def use(self, input: str) -> str:
        try:
            answer, url = input.split("|||")
            return f"{answer.strip()}\n\nSource: [{url.strip()}]({url.strip()})"
        except:
            return "Could not format citation correctly."

summarizer = pipeline("summarization")
@tool
class SummarizerTool:
    name = "summarize"
    description = "Summarize a long text into a short paragraph."
    def use(self, input: str) -> str:
        if len(input) < 50:
            return input
        result = summarizer(input, max_length=100, min_length=25, do_sample=False)
        return result[0]['summary_text']

@tool
class PythonTool:
    name = "python"
    description = "Execute Python code to solve math problems."
    def use(self, code: str) -> str:
        try:
            result = str(eval(code, {"__builtins__": {}}))
            return f"Answer: {result}"
        except Exception as e:
            return f"Error: {str(e)}"

@tool
class FallbackTool:
    name = "fallback"
    description = "Handle unanswerable or unclear queries."
    def use(self, _: str) -> str:
        return "I'm sorry, I couldn't find the answer to your question. Could you rephrase or try something else?"

# --- Basic Agent Definition ---
class BasicAgent:
    def __init__(self):
        tools = [WebSearchTool(), CiteTool(), SummarizerTool(), PythonTool(), FallbackTool()]
        self.agent = Agent(
            tools=tools,
            system_prompt="""
You are Smart Answering Agent v3.
Answer questions factually, concisely, and cite sources when available.
Route to the correct tool for factual, math, or summarization queries.
If you don’t know the answer, respond gracefully using the fallback tool.
Ensure output format is friendly for the GAIA evaluation.
"""
        )
    def __call__(self, question: str) -> str:
        return self.agent.run(question)

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

def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")
    if profile:
        username = profile.username
    else:
        return "Please Login to Hugging Face with the button.", None

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

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

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

    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 = []

    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 produce any answers to submit.", pd.DataFrame(results_log)

    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }

    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"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("# Smart Agent Evaluation Runner")
    gr.Markdown("""
    **Instructions:**
    1. Login to your HF account using the button.
    2. Click 'Run Evaluation & Submit All Answers' to test your agent.
    """)

    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

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

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