"""Maximum message termination condition."""

from typing import Optional, Sequence

from agentbyte.messages import Message
from agentbyte.types import StopMessage

from .base import BaseTermination


class MaxMessageTermination(BaseTermination):
    """
    Terminates when maximum message count is reached.

    Counts all messages (user, assistant, system, tool) streamed
    during orchestration.

    Example:
        termination = MaxMessageTermination(max_messages=10)
        orchestrator = RoundRobinOrchestrator(
            agents=[agent1, agent2],
            termination=termination
        )
    """

    def __init__(self, max_messages: int):
        super().__init__()
        if max_messages <= 0:
            raise ValueError("max_messages must be > 0")
        self.max_messages = max_messages
        self.message_count = 0

    def check(self, new_messages: Sequence[Message]) -> Optional[StopMessage]:
        """Check if message limit is exceeded."""
        self.message_count += len(new_messages)

        if self.message_count >= self.max_messages:
            return self._set_termination(
                f"Maximum messages reached ({self.message_count}/{self.max_messages})",
                {
                    "message_count": self.message_count,
                    "max_messages": self.max_messages,
                },
            )

        return None

    def reset(self) -> None:
        """Reset message counter."""
        super().reset()
        self.message_count = 0
