from __future__ import annotations

import uuid
import asyncio
from pydantic import BaseModel
from typing import List

from autogen_core import (
    SingleThreadedAgentRuntime,
    TypeSubscription,
)
from autogen_core.models import AssistantMessage, LLMMessage
from autogen_core import (
    MessageContext,
    RoutedAgent,
    TopicId,
    AgentId,
    message_handler,
)


INTERACTION_AGENT_TOPIC_TYPE = 'INTERACTION_AGENT_TOPIC_TYPE'


class OutputMessage(BaseModel):
    source: str
    context: List[LLMMessage]


class InteractionAgent(RoutedAgent):
    def __init__(self,
                 description):
        super(InteractionAgent, self).__init__(description=description)

    @message_handler
    async def handle_input(self, message: OutputMessage, ctx: MessageContext) -> None:
        """
        处理规划型任务
        Returns:

        """

        user_input = "test_task"
        publish_msg = message
        await self.send_message(publish_msg, recipient=AgentId('not_exists_agent_name', 'Test'))


def interaction_agent_factory() -> InteractionAgent:
    return InteractionAgent(
        description='A agent to be used to interact with human',
    )


async def main():

    runtime = SingleThreadedAgentRuntime()

    interaction_agent_type = await InteractionAgent.register(runtime=runtime, type=INTERACTION_AGENT_TOPIC_TYPE,
                                                             factory=interaction_agent_factory)

    await runtime.add_subscription(
        TypeSubscription(topic_type=INTERACTION_AGENT_TOPIC_TYPE, agent_type=interaction_agent_type))

    runtime.start()
    session_id = str(uuid.uuid4())

    start_msg = AssistantMessage(content='Hello, how can i help you.', source='assistant')
    start_message = OutputMessage(source='system',
                                  context=[start_msg])

    await runtime.publish_message(start_message, topic_id=TopicId(INTERACTION_AGENT_TOPIC_TYPE, source=session_id))
    await runtime.stop_when_idle()
    await runtime.close()

if __name__ == '__main__':
    asyncio.run(main())

    # loop = asyncio.get_event_loop()
    #
    # loop.run_until_complete(main())