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

from smolagents import CodeAgent, OpenAIServerModel
from smolagents.tools.web_search import WebSearchTool

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

class RetryableWebSearchTool(WebSearchTool):
    def run(self, query: str) -> str:
        for attempt in range(3):
            try:
                return super().run(query)
            except Exception as e:
                if "rate" in str(e).lower():
                    print(f"[WebSearch] Rate limit, retry {attempt+1}/3")
                    time.sleep(2 * (attempt + 1))
                else:
                    raise
        raise RuntimeError("Web search failed after retries")

class SmartGAIAAgent:
    def __init__(self):
        key = os.getenv("OPENAI_API_KEY")
        if not key:
            raise ValueError("Missing OPENAI_API_KEY")
        model = OpenAIServerModel(model_id="gpt-4", api_key=key)
        self.agent = CodeAgent(
            model=model,
            tools=[RetryableWebSearchTool()],
            add_base_tools=True
        )

    def __call__(self, question: str) -> str:
        question = question[:MAX_QUESTION_LENGTH]
        try:
            return self.agent.run(question).strip()
        except Exception as e:
            print("Agent error:", e)
            return "error"

def run_and_submit_all(profile: gr.OAuthProfile | None):
    username = profile.username if profile else None
    if not username:
        return "Please Login to Hugging Face", None

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

    response = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
    response.raise_for_status()
    questions = response.json()

    payload = []
    logs = []

    for item in questions:
        tid = item.get("task_id")
        q = item.get("question", "")
        if not tid or not q or len(q) > MAX_QUESTION_LENGTH:
            continue
        if any(skip in q.lower() for skip in ['.mp3', '.wav', '.jpg', '.png', 'youtube', 'video', 'watch', 'listen']):
            continue
        answer = agent(q)
        payload.append({"task_id": tid, "submitted_answer": answer})
        logs.append({"Task ID": tid, "Question": q, "Submitted Answer": answer})

    if not payload:
        return "No valid questions to submit.", pd.DataFrame(logs)

    submission = {
        "username": username,
        "agent_code": f"https://huggingface.co/spaces/{os.getenv('SPACE_ID')}/tree/main",
        "answers": payload,
    }

    resp = requests.post(f"{DEFAULT_API_URL}/submit", json=submission, timeout=30)
    resp.raise_for_status()
    result = resp.json()
    status = f"Score: {result.get('score')}% ({result.get('correct_count')}/{result.get('total_attempted')})"
    return status, pd.DataFrame(logs)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent")
    gr.LoginButton()
    run_btn = gr.Button("Run & Submit")
    status = gr.Textbox(lines=5)
    table = gr.DataFrame()
    run_btn.click(run_and_submit_all, outputs=[status, table])

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