from autogen_agentchat.agents import (
    AssistantAgent,
    MessageFilterAgent,
    MessageFilterConfig,
    PerSourceFilter,
)
from autogen_agentchat.teams import (
    DiGraphBuilder,
    GraphFlow,
)
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.tools import FunctionTool

from duckduckgo_search import DDGS


model_client = OpenAIChatCompletionClient(model="gpt-4o")

# Agents
generator = AssistantAgent(
    "generator",
    model_client=model_client,
    system_message="Generate a list of creative ideas.",
)
reviewer = AssistantAgent(
    "reviewer",
    model_client=model_client,
    system_message="Review ideas and provide feedbacks, or just 'APPROVE' for final approval.",
)
summarizer_core = AssistantAgent(
    "summary",
    model_client=model_client,
    system_message="Summarize the user request and the final feedback.",
)

# additional agent
# basic web surfer agent to search the web for information using duckduckgo_search


def web_search(search: str) -> str:
    return DDGS().text(search, max_results=5)


web_search_tool = FunctionTool(
    web_search, description="Search the web for information."
)

web_surfer = AssistantAgent(
    "web_surfer",
    model_client,
    system_message="You are a web surfer agent. Your task is to search the web for information based on the instructions provided by the planner agent. You will gather information and present it to the fact-finder agent. If you need more specific details, ask the planner agent for clarification.",
    tools=[web_search_tool],
)

# Filtered summarizer
filtered_summarizer = MessageFilterAgent(
    name="summary",
    wrapped_agent=summarizer_core,
    filter=MessageFilterConfig(
        per_source=[
            PerSourceFilter(source="user", position="first", count=1),
            PerSourceFilter(source="reviewer", position="last", count=1),
        ]
    ),
)


async def main() -> None:
    builder = DiGraphBuilder()
    builder.add_node(generator).add_node(reviewer).add_node(
        filtered_summarizer
    ).add_node(web_surfer)  # add new web surfer agent

    # add web surfer node between generator and reviewer
    builder.add_edge(generator, web_surfer)
    # builder.add_edge(generator, reviewer) # original connection
    builder.add_edge(web_surfer, reviewer)
    builder.add_edge(
        reviewer,
        filtered_summarizer,
        condition=lambda msg: "APPROVE" in msg.to_model_text(),
    )
    builder.add_edge(
        reviewer,
        web_surfer,  # was generator
        condition=lambda msg: "APPROVE" not in msg.to_model_text(),
    )
    builder.set_entry_point(generator)
    graph = builder.build()
    flow = GraphFlow(
        participants=builder.get_participants(),
        graph=graph,
    )

    await Console(flow.run_stream(task="Brainstorm ways to reduce plastic waste."))


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())