import os
import requests
from markdownify import markdownify as md
from requests.exceptions import RequestException
import re
from transformers.agents import (
    ReactCodeAgent,
    ReactJsonAgent,
    HfApiEngine,
    ManagedAgent,
    stream_to_gradio,
    Tool,
)
from transformers.agents.search import DuckDuckGoSearchTool
import gradio as gr
import datetime
from huggingface_hub import login

# Log in to Hugging Face
hf_token = os.getenv("hf_token")
login(hf_token)
model = "meta-llama/Meta-Llama-3.1-70B-Instruct"

# Define the VisitWebpageTool
class VisitWebpageTool(Tool):
    name = "visit_webpage"
    description = "Visits a webpage at the given url and returns its content as a markdown string."
    inputs = {
        "url": {
            "type": "string",
            "description": "The url of the webpage to visit.",
        }
    }
    output_type = "string"

    def forward(self, url: str) -> str:
        try:
            response = requests.get(url)
            response.raise_for_status()
            markdown_content = md(response.text).strip()
            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
            return markdown_content
        except RequestException as e:
            return f"Error fetching the webpage: {str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

visit_page_tool = VisitWebpageTool()

# Set up the multi-agent system
llm_engine = HfApiEngine(model)

# Agent to clarify specifics
clarify_agent = ReactJsonAgent(
    tools=[],
    llm_engine=llm_engine,
    max_iterations=20,
)

# Agent to plan the project
plan_agent = ReactJsonAgent(
    tools=[],
    llm_engine=llm_engine,
    max_iterations=20,
)

# Agent to break the plan into tasks
task_agent = ReactJsonAgent(
    tools=[],
    llm_engine=llm_engine,
    max_iterations=20,
)

# Agent to execute tasks
execute_agent = ReactCodeAgent(
    tools=[],
    llm_engine=llm_engine,
    max_iterations=20,
    additional_authorized_imports=['requests', 'bs4']
)

# Agent to review code
review_agent = ReactCodeAgent(
    tools=[],
    llm_engine=llm_engine,
    max_iterations=20,
    additional_authorized_imports=['requests', 'bs4', 'transformers']
    )

# Agent to revise code
revise_agent = ReactCodeAgent(
    tools=[],
    llm_engine=llm_engine,
    max_iterations=20,
)

# Managed agents
managed_clarify_agent = ManagedAgent(
    agent=clarify_agent,
    name="clarify_agent",
    description="Clarifies project specifics if not given in the project description.",
)

managed_plan_agent = ManagedAgent(
    agent=plan_agent,
    name="plan_agent",
    description="Plans the project.",
)

managed_task_agent = ManagedAgent(
    agent=task_agent,
    name="task_agent",
    description="Breaks the plan into tasks.",
)

managed_execute_agent = ManagedAgent(
    agent=execute_agent,
    name="execute_agent",
    description="Executes the tasks.",
)

managed_review_agent = ManagedAgent(
    agent=review_agent,
    name="review_agent",
    description="Reviews the code written in completion of tasks.",
)

managed_revise_agent = ManagedAgent(
    agent=revise_agent,
    name="revise_agent",
    description="Revises the code if needed.",
)

# Manager agent
manager_agent = ReactCodeAgent(
    tools=[],
    llm_engine=llm_engine,
    managed_agents=[
        managed_clarify_agent,
        managed_plan_agent,
        managed_task_agent,
        managed_execute_agent,
        managed_review_agent,
        managed_revise_agent,
    ],
    additional_authorized_imports=["time", "datetime"],
)

# Implement the Gradio interface
def interact_with_agent(task):
    messages = []
    messages.append(gr.ChatMessage(role="user", content=task))
    yield messages
    try:
        for msg in stream_to_gradio(manager_agent, task):
            messages.append(msg)
            yield messages + [
                gr.ChatMessage(role="assistant", content="⏳ Task not finished yet!")
            ]
    except Exception as e:
        messages.append(gr.ChatMessage(role="assistant", content=f"Error: {str(e)}"))
    yield messages

with gr.Blocks() as demo:
    gr.Markdown("# Multi-agent Software Team")
    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="Project Description", value="Create a simple web app that allows users to upload images and apply filters.")
    submit = gr.Button("Start Project Development")
    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(share=True)