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

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

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

# Reliable search tool with retry
class ReliableWebSearch(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("WebSearchTool failed after retries")

# Main agent using GPT-4 & tools
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(
            tools=[ReliableWebSearch()],
            model=model,
            add_base_tools=True
        )

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

# Fetch, filter, run, and submit answers
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

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

    payload, logs = [], []
    skip_kw = ['.mp3','.wav','.png','.jpg','youtube','video','watch','listen']
    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 or any(k in q.lower() for k in skip_kw):
            continue
        ans = agent(q)
        payload.append({"task_id": tid, "submitted_answer": ans})
        logs.append({"Task ID": tid, "Question": q, "Submitted Answer": ans})

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

    sub = {
        "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=sub, timeout=60)
    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)