import requests
from smolagents import ToolCallingAgent, TransformersModel, Tool, DuckDuckGoSearchTool

# ----------------------------
# MODEL CONFIGURATION
# ----------------------------
model = TransformersModel(
    model_id="meta-llama/Llama-3.2-3B-Instruct"
)

# ----------------------------
# CUSTOM WEB SEARCH TOOL
# ----------------------------
class WebSearchTool(Tool):
    def __init__(self):
        super().__init__()
        self.name = "web_search"
        self.description = "Search the web using DuckDuckGo"
        self._ddg = DuckDuckGoSearchTool(max_results=5, rate_limit=2.0)

    def __call__(self, query: str):
        return self._ddg(query=query)

# ----------------------------
# AGENT
# ----------------------------
agent = ToolCallingAgent(
    model=model,
    tools=[WebSearchTool()],
    max_steps=8
)

# ----------------------------
# GAIA API ENDPOINTS
# ----------------------------
BASE_URL = "https://agents-course-unit4-scoring.hf.space"

def get_questions():
    r = requests.get(f"{BASE_URL}/questions")
    r.raise_for_status()
    return r.json()

def get_random_question():
    r = requests.get(f"{BASE_URL}/random-question")
    r.raise_for_status()
    return r.json()

def submit_answers(username, agent_code, answers):
    payload = {"username": username, "agent_code": agent_code, "answers": answers}
    r = requests.post(f"{BASE_URL}/submit", json=payload)
    r.raise_for_status()
    return r.json()

# ----------------------------
# ANSWER GENERATION
# ----------------------------
def generate_answer(question):
    prompt = f"""
Answer this question accurately and concisely.
Do not include reasoning or explanations.
Question: {question}
"""
    try:
        return agent.run(prompt).strip()
    except Exception as e:
        return f"Error: {e}"

# ----------------------------
# MAIN SCRIPT
# ----------------------------
if __name__ == "__main__":
    # Replace with your Hugging Face username and Space URL
    username = "<YOUR_HF_USERNAME>"
    agent_code = "https://huggingface.co/spaces/<YOUR_HF_USERNAME>/smolagents-final/tree/main"

    # Retrieve all GAIA questions
    questions = get_questions()
    print(f"Retrieved {len(questions)} GAIA questions.\n")

    # Generate answers
    answers = []
    for q in questions:
        print(f"[{q['task_id']}] {q['question']}")
        ans = generate_answer(q["question"])
        print(f" → {ans}")
        answers.append({"task_id": q["task_id"], "submitted_answer": ans})

    # Submit answers
    print("\nSubmitting answers...")
    result = submit_answers(username, agent_code, answers)
    print(result)