import os
import time
import random
import spaces
import gradio as gr
from transformers import (
    ReactCodeAgent,
    ReactJsonAgent,
    HfApiEngine,
    ManagedAgent,
    stream_to_gradio,
)
from transformers.agents.search import DuckDuckGoSearchTool
from visit_webpage_tool import VisitWebpageTool  # Import the VisitWebpageTool class
from huggingface_hub import login

# Define the model
model = "meta-llama/Meta-Llama-3.1-70B-Instruct"

# Initialize the LLM engine with the Hugging Face token
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    raise ValueError("Hugging Face API token not found. Please set the hf_token environment variable.")

# Authenticate with Hugging Face
login(hf_token)

# Initialize the LLM engine with reduced max_new_tokens
llm_engine = HfApiEngine(model)

# Create the web agent
web_agent = ReactJsonAgent(
    tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
    llm_engine=llm_engine,
    max_iterations=10,
)

# Wrap the web agent into a ManagedAgent
managed_web_agent = ManagedAgent(
    agent=web_agent,
    name="search_agent",
    description="Runs web searches for you. Give it your query as an argument.",
)

# Create the manager agent
manager_agent = ReactCodeAgent(
    tools=[],
    llm_engine=llm_engine,
    managed_agents=[managed_web_agent],
    additional_authorized_imports=["time", "datetime"],
)

# Define a retry mechanism with exponential backoff
def retry_with_backoff(func, retries=5, backoff_factor=2):
    for attempt in range(retries):
        try:
            return func()
        except Exception as e:
            if attempt < retries - 1:
                sleep_time = backoff_factor ** attempt + random.uniform(0, 1)
                time.sleep(sleep_time)
            else:
                raise e

# Define the Gradio interaction function

def interact_with_agent(task):
    messages = []
    messages.append(gr.ChatMessage(role="user", content=task))
    yield messages
    for msg in stream_to_gradio(manager_agent, task):
        messages.append(msg)
        yield messages + [
            gr.ChatMessage(role="assistant", content="⏳ Task not finished yet!")
        ]
    yield messages

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# multi-agent-web-browser")
    gr.Markdown("Gradio space based on the multiagent_web_assistant cookbook https://huggingface.co/learn/cookbook/multiagent_web_assistant")
    text_input = gr.Textbox(lines=1, label="Chat Message", value="How many years ago was Stripe founded?")
    submit = gr.Button("Run web search agent!")
    chatbot = gr.Chatbot(
        label="Agent",
        type="messages",
        avatar_images=(
            None,
            "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",
        ),
    )
    submit.click(interact_with_agent, [text_input], [chatbot])

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