import os
import re
import gradio as gr

from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel

# ============================================================
# API KEY
# ============================================================

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# ============================================================
# MODEL
# ============================================================

model = LiteLLMModel(
    model_id="groq/llama-3.3-70b-versatile",
    api_key=GROQ_API_KEY,
    temperature=0.2,
    max_tokens=512
)

# ============================================================
# TOOLS
# ============================================================

search_tool = DuckDuckGoSearchTool()

agent = CodeAgent(
    tools=[search_tool],
    model=model,
    max_steps=5,
    verbosity_level=1
)

# ============================================================
# CLEAN OUTPUT
# ============================================================

def clean_output(text):
    if not text:
        return ""

    text = str(text).strip()

    text = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
    text = re.sub(r"(?i)(final answer:|answer:)", "", text)

    return text.strip()

# ============================================================
# SOLVER
# ============================================================

def solve(question):

    prompt = f"""
You are solving a GAIA benchmark task.

Question:
{question}

IMPORTANT:
- Return ONLY the final answer
- No explanation
- No markdown
- No extra words
"""

    try:
        result = agent.run(prompt)
        return clean_output(result)

    except Exception as e:
        return f"ERROR: {str(e)}"

# ============================================================
# UI
# ============================================================

demo = gr.Interface(
    fn=solve,
    inputs=gr.Textbox(lines=8, label="Question"),
    outputs="text",
    title="GAIA Agent"
)

demo.launch(server_name="0.0.0.0", server_port=7860)