# Copyright (c) Microsoft. All rights reserved.

import asyncio
import logging
import sys
from collections.abc import Awaitable, Callable

from semantic_kernel.agents.agent import Agent
from semantic_kernel.agents.orchestration.agent_actor_base import ActorBase, AgentActorBase
from semantic_kernel.agents.orchestration.orchestration_base import (
    ChatMessageContent,
    DefaultTypeAlias,
    OrchestrationBase,
    TIn,
    TOut,
)
from semantic_kernel.agents.runtime.core.cancellation_token import CancellationToken
from semantic_kernel.agents.runtime.core.core_runtime import CoreRuntime
from semantic_kernel.agents.runtime.core.message_context import MessageContext
from semantic_kernel.agents.runtime.core.routed_agent import message_handler
from semantic_kernel.agents.runtime.core.topic import TopicId
from semantic_kernel.agents.runtime.in_process.type_subscription import TypeSubscription
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
from semantic_kernel.kernel_pydantic import KernelBaseModel
from semantic_kernel.utils.feature_stage_decorator import experimental

if sys.version_info >= (3, 12):
    from typing import override  # pragma: no cover
else:
    from typing_extensions import override  # pragma: no cover


logger: logging.Logger = logging.getLogger(__name__)


@experimental
class ConcurrentRequestMessage(KernelBaseModel):
    """A request message type for concurrent agents."""

    body: DefaultTypeAlias


@experimental
class ConcurrentResponseMessage(KernelBaseModel):
    """A response message type for concurrent agents."""

    body: ChatMessageContent


@experimental
class ConcurrentAgentActor(AgentActorBase):
    """A agent actor for concurrent agents that process tasks."""

    def __init__(
        self,
        agent: Agent,
        internal_topic_type: str,
        collection_agent_type: str,
        exception_callback: Callable[[BaseException], None],
        agent_response_callback: Callable[[DefaultTypeAlias], Awaitable[None] | None] | None = None,
        streaming_agent_response_callback: Callable[[StreamingChatMessageContent, bool], Awaitable[None] | None]
        | None = None,
    ) -> None:
        """Initialize the agent actor.

        Args:
            agent: The agent to be executed.
            internal_topic_type: The internal topic type for the actor.
            collection_agent_type: The collection agent type for the actor.
            exception_callback: A callback function to handle exceptions.
            agent_response_callback: A callback function to handle the full response from the agent.
            streaming_agent_response_callback: A callback function to handle streaming responses from the agent.
        """
        self._collection_agent_type = collection_agent_type
        super().__init__(
            agent=agent,
            internal_topic_type=internal_topic_type,
            exception_callback=exception_callback,
            agent_response_callback=agent_response_callback,
            streaming_agent_response_callback=streaming_agent_response_callback,
        )

    @message_handler
    async def _handle_message(self, message: ConcurrentRequestMessage, ctx: MessageContext) -> None:
        """Handle a message."""
        logger.debug(f"Concurrent actor (Actor ID: {self.id}; Agent name: {self._agent.name}) started processing...")

        response = await self._invoke_agent(additional_messages=message.body)

        logger.debug(f"Concurrent actor (Actor ID: {self.id}; Agent name: {self._agent.name}) finished processing.")

        target_actor_id = await self.runtime.get(self._collection_agent_type)
        await self.send_message(
            ConcurrentResponseMessage(body=response),
            target_actor_id,
            cancellation_token=ctx.cancellation_token,
        )


@experimental
class CollectionActor(ActorBase):
    """A agent container for collecting results from concurrent agents."""

    def __init__(
        self,
        description: str,
        expected_answer_count: int,
        exception_callback: Callable[[BaseException], None],
        result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
    ) -> None:
        """Initialize the collection agent container."""
        self._expected_answer_count = expected_answer_count
        self._result_callback = result_callback
        self._results: list[ChatMessageContent] = []
        self._lock = asyncio.Lock()

        super().__init__(description, exception_callback)

    @message_handler
    async def _handle_message(self, message: ConcurrentResponseMessage, _: MessageContext) -> None:
        async with self._lock:
            self._results.append(message.body)

        if len(self._results) == self._expected_answer_count:
            logger.debug(f"Collection actor (Actor ID: {self.id}) finished processing all responses.")
            if self._result_callback:
                await self._result_callback(self._results)


@experimental
class ConcurrentOrchestration(OrchestrationBase[TIn, TOut]):
    """A concurrent multi-agent pattern orchestration."""

    @override
    async def _start(
        self,
        task: DefaultTypeAlias,
        runtime: CoreRuntime,
        internal_topic_type: str,
        cancellation_token: CancellationToken,
    ) -> None:
        """Start the concurrent pattern."""
        await runtime.publish_message(
            ConcurrentRequestMessage(body=task),
            TopicId(internal_topic_type, self.__class__.__name__),
            cancellation_token=cancellation_token,
        )

    @override
    async def _prepare(
        self,
        runtime: CoreRuntime,
        internal_topic_type: str,
        exception_callback: Callable[[BaseException], None],
        result_callback: Callable[[DefaultTypeAlias], Awaitable[None]],
    ) -> None:
        """Register the actors and orchestrations with the runtime and add the required subscriptions."""
        await asyncio.gather(*[
            self._register_members(
                runtime,
                internal_topic_type,
                exception_callback,
            ),
            self._register_collection_actor(
                runtime,
                internal_topic_type,
                exception_callback,
                result_callback=result_callback,
            ),
            self._add_subscriptions(
                runtime,
                internal_topic_type,
            ),
        ])

    async def _register_members(
        self,
        runtime: CoreRuntime,
        internal_topic_type: str,
        exception_callback: Callable[[BaseException], None],
    ) -> None:
        """Register the members."""

        async def _internal_helper(agent: Agent) -> None:
            await ConcurrentAgentActor.register(
                runtime,
                self._get_agent_actor_type(agent, internal_topic_type),
                lambda agent=agent: ConcurrentAgentActor(  # type: ignore[misc]
                    agent,
                    internal_topic_type,
                    self._get_collection_actor_type(internal_topic_type),
                    exception_callback,
                    agent_response_callback=self._agent_response_callback,
                    streaming_agent_response_callback=self._streaming_agent_response_callback,
                ),
            )

        await asyncio.gather(*[_internal_helper(agent) for agent in self._members])

    async def _register_collection_actor(
        self,
        runtime: CoreRuntime,
        internal_topic_type: str,
        exception_callback: Callable[[BaseException], None],
        result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
    ) -> None:
        await CollectionActor.register(
            runtime,
            self._get_collection_actor_type(internal_topic_type),
            lambda: CollectionActor(
                description="An internal agent that is responsible for collection results",
                expected_answer_count=len(self._members),
                exception_callback=exception_callback,
                result_callback=result_callback,
            ),
        )

    async def _add_subscriptions(
        self,
        runtime: CoreRuntime,
        internal_topic_type: str,
    ) -> None:
        await asyncio.gather(*[
            runtime.add_subscription(
                TypeSubscription(
                    internal_topic_type,
                    self._get_agent_actor_type(agent, internal_topic_type),
                )
            )
            for agent in self._members
        ])

    def _get_agent_actor_type(self, worker: Agent, internal_topic_type: str) -> str:
        """Get the container type for an agent.

        The type is appended with the internal topic type to ensure uniqueness in the runtime
        that may be shared by multiple orchestrations.
        """
        return f"{worker.name}_{internal_topic_type}"

    def _get_collection_actor_type(self, internal_topic_type: str) -> str:
        """Get the collection agent type.

        The type is appended with the internal topic type to ensure uniqueness in the runtime
        that may be shared by multiple orchestrations.
        """
        return f"{CollectionActor.__name__}_{internal_topic_type}"