import os
import gradio as gr
import requests
import pandas as pd
from transformers import pipeline
from duckduckgo_search import ddg

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

# --- Tool Definitions ---
class WebSearchTool:
    """A simple open-source web search tool using DuckDuckGo."""
    def search(self, query: str) -> list[str]:
        # Perform a DuckDuckGo search and return top 3 titles
        try:
            results = ddg(query, max_results=3)
            if not results:
                return []
            return [item.get("title", "").strip() for item in results]
        except Exception as e:
            print(f"DuckDuckGo search error: {e}")
            return []

class CalculatorTool:
    """Evaluates simple arithmetic expressions."""
    def calculate(self, expression: str) -> float:
        # WARNING: using eval; in production, use a safe parser
        return eval(expression, {}, {})

# --- Agent Definition ---
class ToolsAgent:
    def __init__(self):
        print("ToolsAgent initialized.")
        # Initialize LLM for general reasoning
        self.llm = pipeline(
            "text-generation", model="gpt2", tokenizer="gpt2", return_full_text=False
        )
        # Initialize tools
        self.tools = {
            "search": WebSearchTool(),
            "calc": CalculatorTool(),
        }

    def __call__(self, question: str) -> str:
        print(f"Agent received question: {question}")
        q_lower = question.lower().strip()
        # 1) If calculation question
        if any(op in q_lower for op in ["+", "-", "*", "/"]):
            try:
                result = self.tools["calc"].calculate(q_lower)
                return str(result)
            except Exception as e:
                print(f"Calc error: {e}")
        # 2) If search instruction
        if q_lower.startswith("search") or "find" in q_lower:
            try:
                titles = self.tools["search"].search(question)
                # Return only the first word of the top title
                return titles[0].split()[0] if titles else ""
            except Exception as e:
                print(f"Search error: {e}")
        # 3) Fallback to LLM
        try:
            out = self.llm(question, max_length=50, num_return_sequences=1)[0]["generated_text"].strip()
            # Return only the first token (single word/number)
            return out.split()[0]
        except Exception as e:
            print(f"LLM error: {e}")
            return ""

# --- Evaluation and Submission Logic ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")
    if not profile:
        return "Please Login to Hugging Face with the button.", None
    username = profile.username.strip()

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

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

    try:
        resp = requests.get(questions_url, timeout=15)
        resp.raise_for_status()
        qs = resp.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    answers, log = [], []
    for item in qs:
        tid = item.get("task_id")
        q = item.get("question")
        if not tid or q is None:
            continue
        ans = agent(q)
        answers.append({"task_id": tid, "submitted_answer": ans})
        log.append({"Task ID": tid, "Question": q, "Submitted Answer": ans})

    if not answers:
        return "No answers generated.", pd.DataFrame(log)

    payload = {
        "username": username,
        "agent_code": f"https://huggingface.co/spaces/{space_id}/tree/main",
        "answers": answers
    }
    try:
        res = requests.post(submit_url, json=payload, timeout=60)
        res.raise_for_status()
        data = res.json()
        status = (
            f"Submission Successful! User: {data.get('username')} "
            f"Score: {data.get('score', 'N/A')}%"
        )
    except Exception as e:
        status = f"Submission Failed: {e}"
    return status, pd.DataFrame(log)

# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("## Agents Course: DuckDuckGo Search + Calculator Agent")
    gr.LoginButton()
    run_btn = gr.Button("Run Evaluation & Submit All Answers")
    out_txt = gr.Textbox(label="Result", lines=3)
    out_df = gr.DataFrame(label="Log", wrap=True)
    run_btn.click(fn=run_and_submit_all, outputs=[out_txt, out_df])

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