# Copyright (c) Microsoft. All rights reserved.

"""Base class for group chat orchestrators that manages conversation flow and participant selection."""

import asyncio
import inspect
import logging
import sys
from abc import ABC
from collections import OrderedDict
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import Any, ClassVar, TypeAlias

from agent_framework._types import AgentResponse, AgentResponseUpdate, Message
from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._events import WorkflowEvent
from agent_framework._workflows._executor import Executor, handler
from agent_framework._workflows._workflow_context import WorkflowContext
from typing_extensions import Never

from ._orchestration_request_info import AgentApprovalExecutor

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


logger = logging.getLogger(__name__)


@dataclass
class GroupChatRequestMessage:
    """Request envelope sent from the orchestrator to a participant."""

    additional_instruction: str | None = None
    metadata: dict[str, Any] | None = None


@dataclass
class GroupChatParticipantMessage:
    """Message envelop containing messages generated by a participant.

    This message envelope is used to broadcast messages from one participant
    to other participants in the group chat to keep them synchronized.
    """

    messages: list[Message]


@dataclass
class GroupChatResponseMessage:
    """Response envelope emitted by participants back to the orchestrator."""

    message: Message


TerminationCondition: TypeAlias = Callable[[list[Message]], bool | Awaitable[bool]]
GroupChatWorkflowContextOutT: TypeAlias = AgentExecutorRequest | GroupChatRequestMessage | GroupChatParticipantMessage


# region Group chat events


@dataclass
class GroupChatRequestSentEvent:
    """Data payload for group_chat request sent events."""

    round_index: int
    participant_name: str


@dataclass
class GroupChatResponseReceivedEvent:
    """Data payload for group_chat response received events."""

    round_index: int
    participant_name: str


# endregion


# region Participant registry
class ParticipantRegistry:
    """Simple registry for tracking group chat participants and their types and other properties."""

    EMPTY_DESCRIPTION_PLACEHOLDER: ClassVar[str] = (
        "<no description, use name to identify the purpose of this participant>"
    )

    def __init__(self, participants: Sequence[Executor]) -> None:
        """Initialize the registry and validate participant IDs.

        Args:
            participants: List of executors (agents or custom executors) to register
        Raises:
            ValueError: If there are duplicate or conflicting participant IDs
        """
        self._agents: set[str] = set()
        self._participants: OrderedDict[str, str] = OrderedDict()
        self._resolve_participants(participants)

    def _resolve_participants(self, participants: Sequence[Executor]) -> None:
        """Register participants and validate IDs."""
        for participant in participants:
            if participant.id in self._participants:
                raise ValueError(f"Participant ID conflict: '{participant.id}' registered as both agent and executor.")

            if isinstance(participant, AgentExecutor | AgentApprovalExecutor):
                self._agents.add(participant.id)
                self._participants[participant.id] = participant.description or self.EMPTY_DESCRIPTION_PLACEHOLDER
            else:
                self._participants[participant.id] = self.EMPTY_DESCRIPTION_PLACEHOLDER

    def is_agent(self, name: str) -> bool:
        """Check if a participant is an agent (vs custom executor)."""
        return name in self._agents

    @property
    def participants(self) -> OrderedDict[str, str]:
        """Get all registered participant names and descriptions in an ordered dictionary."""
        return self._participants


# endregion


class BaseGroupChatOrchestrator(Executor, ABC):
    """Abstract base class for group chat orchestrators.

    Provides shared functionality for participant registration, routing,
    and round limit checking that is common across all group chat patterns.

    Subclasses must implement pattern-specific orchestration logic while
    inheriting the common participant management infrastructure.
    """

    TERMINATION_CONDITION_MET_MESSAGE: ClassVar[str] = "The group chat has reached its termination condition."
    MAX_ROUNDS_MET_MESSAGE: ClassVar[str] = "The group chat has reached the maximum number of rounds."

    def __init__(
        self,
        id: str,
        participant_registry: ParticipantRegistry,
        *,
        name: str | None = None,
        max_rounds: int | None = None,
        termination_condition: TerminationCondition | None = None,
    ) -> None:
        """Initialize base orchestrator.

        Args:
            id: Unique identifier for this orchestrator executor
            participant_registry: Registry of group chat participants that tracks their types (agents
                vs custom executors)
            name: Optional display name for orchestrator messages
            max_rounds: Optional maximum number of conversation rounds.
                Must be equal to or greater than 1 if set. Number smaller than 1 will be coerced to 1.
            termination_condition: Optional callable to determine conversation termination
        """
        super().__init__(id)
        self._name = name or id
        self._max_rounds = max(1, max_rounds) if max_rounds is not None else None
        self._termination_condition = termination_condition
        self._round_index: int = 0
        self._participant_registry = participant_registry
        # Shared conversation state management
        self._full_conversation: list[Message] = []

    # region Handlers

    @handler
    async def handle_str(
        self,
        task: str,
        ctx: WorkflowContext[GroupChatWorkflowContextOutT, list[Message]],
    ) -> None:
        """Handler for string input as workflow entry point.

        Wraps the string in a USER role Message and delegates to _handle_task_message.

        Args:
            task: Plain text task description from user
            ctx: Workflow context

        Usage:
            workflow.run("Write a blog post about AI agents")
        """
        await self._handle_messages([Message(role="user", contents=[task])], ctx)

    @handler
    async def handle_message(
        self,
        task: Message,
        ctx: WorkflowContext[GroupChatWorkflowContextOutT, list[Message]],
    ) -> None:
        """Handler for single Message input as workflow entry point.

        Wraps the message in a list and delegates to _handle_task_message.

        Args:
            task: Message from user
            ctx: Workflow context

        Usage:
            workflow.run(Message(role="user", contents=["Write a blog post about AI agents"]))
        """
        await self._handle_messages([task], ctx)

    @handler
    async def handle_messages(
        self,
        task: list[Message],
        ctx: WorkflowContext[GroupChatWorkflowContextOutT, list[Message]],
    ) -> None:
        """Handler for list of ChatMessages as workflow entry point.

        Delegates to _handle_task_message.

        Args:
            task: List of ChatMessages from user
            ctx: Workflow context
        Usage:
            workflow.run([
                Message(role="user", contents=["Write a blog post about AI agents"]),
                Message(role="user", contents=["Make it engaging and informative."])
            ])
        """
        if not task:
            raise ValueError("At least one Message is required to start the group chat workflow.")
        await self._handle_messages(task, ctx)

    @handler
    async def handle_participant_response(
        self,
        response: AgentExecutorResponse | GroupChatResponseMessage,
        ctx: WorkflowContext[GroupChatWorkflowContextOutT, list[Message]],
    ) -> None:
        """Handler for participant responses.

        This method can be overridden by subclasses if specific response handling is needed.

        Args:
            response: Response from a participant
            ctx: Workflow context
        """
        await ctx.add_event(
            WorkflowEvent(
                "group_chat",
                data=GroupChatResponseReceivedEvent(
                    round_index=self._round_index,
                    participant_name=ctx.source_executor_ids[0] if ctx.source_executor_ids else "unknown",
                ),
            )
        )
        await self._handle_response(response, ctx)

    # endregion

    # region Handler methods subclasses must implement

    async def _handle_messages(
        self,
        messages: list[Message],
        ctx: WorkflowContext[GroupChatWorkflowContextOutT, list[Message]],
    ) -> None:
        """Handle task messages from users as workflow entry point.

        Subclasses must implement this method to define pattern-specific orchestration logic.

        Args:
            messages: Task messages from user
            ctx: Workflow context
        """
        raise NotImplementedError("_handle_messages must be implemented by subclasses.")

    async def _handle_response(
        self,
        response: AgentExecutorResponse | GroupChatResponseMessage,
        ctx: WorkflowContext[GroupChatWorkflowContextOutT, list[Message]],
    ) -> None:
        """Handle a participant response.

        Subclasses must implement this method to define pattern-specific response handling logic.

        Args:
            response: Response from a participant
            ctx: Workflow context
        """
        raise NotImplementedError("_handle_response must be implemented by subclasses.")

    # endregion

    # Conversation state management (shared across all patterns)

    def _append_messages(self, messages: Sequence[Message]) -> None:
        """Append messages to the conversation history.

        Args:
            messages: Messages to append
        """
        self._full_conversation.extend(messages)

    def _get_conversation(self) -> list[Message]:
        """Get a copy of the current conversation.

        Returns:
            Cloned conversation list
        """
        return list(self._full_conversation)

    def _process_participant_response(
        self, response: AgentExecutorResponse | GroupChatResponseMessage
    ) -> list[Message]:
        """Extract Message from participant response.

        Args:
            response: Response from participant
        Returns:
            List of ChatMessages extracted from the response
        """
        if isinstance(response, AgentExecutorResponse):
            return response.agent_response.messages
        if isinstance(response, GroupChatResponseMessage):
            return [response.message]
        raise TypeError(f"Unsupported response type: {type(response)}")

    def _clear_conversation(self) -> None:
        """Clear the conversation history."""
        self._full_conversation.clear()

    def _increment_round(self) -> None:
        """Increment the round counter."""
        self._round_index += 1

    async def _check_termination(self) -> bool:
        """Check if conversation should terminate based on termination condition.

        Supports both synchronous and asynchronous termination conditions.

        Returns:
            True if termination condition met, False otherwise
        """
        if self._termination_condition is None:
            return False

        result = self._termination_condition(self._get_conversation())
        if inspect.isawaitable(result):
            result = await result
        return result

    async def _check_terminate_and_yield(
        self, ctx: WorkflowContext[Never, AgentResponse | AgentResponseUpdate]
    ) -> bool:
        """Check termination conditions and yield the completion message if met.

        Args:
            ctx: Workflow context for yielding output

        Returns:
            True if termination condition met and output yielded, False otherwise
        """
        terminate = await self._check_termination()
        if terminate:
            completion_message = self._create_completion_message(self.TERMINATION_CONDITION_MET_MESSAGE)
            self._append_messages([completion_message])
            await self._yield_completion(ctx, completion_message)
            return True

        return False

    async def _yield_completion(
        self,
        ctx: WorkflowContext[Never, AgentResponse | AgentResponseUpdate],
        completion_message: Message,
    ) -> None:
        """Yield a synthesized terminal completion message in the right shape for the run mode.

        Mode-aware to mirror ``AgentExecutor`` semantics:
        - Streaming (``ctx.is_streaming()``): yield a single ``AgentResponseUpdate`` so the
          ``output`` event stream stays uniformly per-chunk.
        - Non-streaming: yield the full ``AgentResponse``.
        """
        if ctx.is_streaming():
            await ctx.yield_output(
                AgentResponseUpdate(
                    contents=list(completion_message.contents),
                    role=completion_message.role,
                    author_name=completion_message.author_name,
                    message_id=completion_message.message_id,
                )
            )
        else:
            await ctx.yield_output(AgentResponse(messages=[completion_message]))

    def _create_completion_message(self, message: str) -> Message:
        """Create a standardized completion message.

        Args:
            message: Completion text

        Returns:
            Message with completion content
        """
        return Message(role="assistant", contents=[message], author_name=self._name)

    # Participant routing (shared across all patterns)

    async def _broadcast_messages_to_participants(
        self,
        messages: list[Message],
        ctx: WorkflowContext[AgentExecutorRequest | GroupChatParticipantMessage],
        participants: Sequence[str] | None = None,
    ) -> None:
        """Broadcast messages to participants.

        This method sends the given messages to all registered participants
        or a specified subset. This acts as a message broadcast mechanism for
        participants in the group chat to stay synchronized.

        Args:
            messages: Messages to send
            ctx: Workflow context for message broadcasting
            participants: Optional list of participant names to route to.
                          If None, routes to all registered participants.
        """
        target_participants = (
            participants if participants is not None else list(self._participant_registry.participants)
        )

        async def _send_messages(target: str) -> None:
            if self._participant_registry.is_agent(target):
                # Send messages without requesting a response
                await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=False), target_id=target)
            else:
                # Send messages wrapped in GroupChatParticipantMessage
                await ctx.send_message(GroupChatParticipantMessage(messages=messages), target_id=target)

        await asyncio.gather(*[_send_messages(p) for p in target_participants])

    async def _send_request_to_participant(
        self,
        target: str,
        ctx: WorkflowContext[AgentExecutorRequest | GroupChatRequestMessage],
        *,
        additional_instruction: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Send a request to a participant.

        This method handles the dual envelope pattern:
        - AgentExecutors receive AgentExecutorRequest (messages only)
        - Custom executors receive GroupChatRequestMessage (full context)

        Args:
            target: Name of the participant to route to
            ctx: Workflow context for message routing
            additional_instruction: Optional additional instruction for the participant.
                This can be used to provide guidance to steer the participant's response.
            metadata: Optional metadata dict

        Raises:
            ValueError: If participant is not registered
        """
        if self._participant_registry.is_agent(target):
            # AgentExecutors receive simple message list
            messages: list[Message] = []
            if additional_instruction:
                messages.append(Message(role="user", contents=[additional_instruction]))
            request = AgentExecutorRequest(messages=messages, should_respond=True)
            await ctx.send_message(request, target_id=target)
            await ctx.add_event(
                WorkflowEvent(
                    "group_chat",
                    data=GroupChatRequestSentEvent(
                        round_index=self._round_index,
                        participant_name=target,
                    ),
                )
            )
        else:
            # Custom executors receive full context envelope
            request = GroupChatRequestMessage(additional_instruction=additional_instruction, metadata=metadata)  # type: ignore[assignment]
            await ctx.send_message(request, target_id=target)
            await ctx.add_event(
                WorkflowEvent(
                    "group_chat",
                    data=GroupChatRequestSentEvent(
                        round_index=self._round_index,
                        participant_name=target,
                    ),
                )
            )

    # Round limit enforcement (shared across all patterns)

    def _check_round_limit(self) -> bool:
        """Check if round limit has been reached.

        Uses instance variables _round_index and _max_rounds.

        Returns:
            True if limit reached, False otherwise
        """
        if self._max_rounds is None:
            return False

        if self._round_index >= self._max_rounds:
            logger.warning(
                "%s reached max_rounds=%s; forcing completion.",
                self.__class__.__name__,
                self._max_rounds,
            )
            return True

        return False

    async def _check_round_limit_and_yield(
        self, ctx: WorkflowContext[Never, AgentResponse | AgentResponseUpdate]
    ) -> bool:
        """Check round limit and yield the max-rounds completion message if reached.

        Args:
            ctx: Workflow context for yielding output

        Returns:
            True if round limit reached and output yielded, False otherwise
        """
        reach_max_rounds = self._check_round_limit()
        if reach_max_rounds:
            completion_message = self._create_completion_message(self.MAX_ROUNDS_MET_MESSAGE)
            self._append_messages([completion_message])
            await self._yield_completion(ctx, completion_message)
            return True

        return False

    # State persistence (shared across all patterns)

    # State persistence (shared across all patterns)

    @override
    async def on_checkpoint_save(self) -> dict[str, Any]:
        """Capture current orchestrator state for checkpointing.

        Default implementation uses OrchestrationState to serialize common state.
        Subclasses can override this method or _snapshot_pattern_metadata() to add pattern-specific data.

        Returns:
            Serialized state dict
        """
        from ._orchestration_state import OrchestrationState

        state = OrchestrationState(
            conversation=list(self._full_conversation),
            round_index=self._round_index,
            orchestrator_name=self._name,
            metadata=self._snapshot_pattern_metadata(),
        )
        return state.to_dict()

    def _snapshot_pattern_metadata(self) -> dict[str, Any]:
        """Serialize pattern-specific state.

        Override this method to add pattern-specific checkpoint data.

        Returns:
            Dict with pattern-specific state (empty by default)
        """
        return {}

    @override
    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
        """Restore orchestrator state from checkpoint.

        Default implementation uses OrchestrationState to deserialize common state.
        Subclasses can override this method or _restore_pattern_metadata() to restore pattern-specific data.

        Args:
            state: Serialized state dict
        """
        from ._orchestration_state import OrchestrationState

        orch_state = OrchestrationState.from_dict(state)
        self._full_conversation = list(orch_state.conversation)
        self._round_index = orch_state.round_index
        self._name = orch_state.orchestrator_name
        self._restore_pattern_metadata(orch_state.metadata)

    def _restore_pattern_metadata(self, metadata: dict[str, Any]) -> None:
        """Restore pattern-specific state.

        Override this method to restore pattern-specific checkpoint data.

        Args:
            metadata: Pattern-specific state dict
        """
        pass
