"""Consecutive agent termination condition."""

from typing import FrozenSet, Optional, Sequence

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

from .base import BaseTermination

_DEFAULT_IGNORED: FrozenSet[str] = frozenset({"orchestrator", "user"})


class ConsecutiveAgentTermination(BaseTermination):
    """
    Terminates when the same agent is selected consecutively N or more times.

    Tracks the ``source`` field across incoming messages.  Messages whose
    source is in ``ignored_sources`` (default: ``{"orchestrator", "user"}``)
    are skipped so that orchestrator bookkeeping messages do not break the
    consecutive run.

    This is a production safety net for ``AIOrchestrator`` and other
    selector-based patterns where the selector might enter an infinite loop
    routing to the same agent.

    Args:
        max_consecutive:  Maximum allowed consecutive turns for one agent
                          before termination fires.
        ignored_sources:  Set of source names to skip when counting.
                          Pass an explicit empty set to track all sources.

    Example:
        # Stop if the same agent is picked 3 times in a row
        termination = ConsecutiveAgentTermination(max_consecutive=3)

        # Combine with a hard message cap
        termination = (
            ConsecutiveAgentTermination(max_consecutive=3)
            | MaxMessageTermination(30)
        )
    """

    def __init__(
        self,
        max_consecutive: int,
        ignored_sources: Optional[Sequence[str]] = None,
    ) -> None:
        super().__init__()
        if max_consecutive <= 0:
            raise ValueError("max_consecutive must be > 0")
        self.max_consecutive = max_consecutive
        self._ignored: FrozenSet[str] = (
            frozenset(ignored_sources) if ignored_sources is not None else _DEFAULT_IGNORED
        )
        self._last_source: Optional[str] = None
        self._consecutive_count: int = 0

    def check(self, new_messages: Sequence[Message]) -> Optional[StopMessage]:
        """Track consecutive same-source turns; terminate when limit is reached."""
        for message in new_messages:
            if not hasattr(message, "source"):
                continue
            source = message.source
            if source in self._ignored:
                continue

            if source == self._last_source:
                self._consecutive_count += 1
            else:
                self._last_source = source
                self._consecutive_count = 1

            if self._consecutive_count >= self.max_consecutive:
                return self._set_termination(
                    f"Agent '{source}' selected {self._consecutive_count} consecutive times "
                    f"(limit: {self.max_consecutive})",
                    {
                        "agent": source,
                        "consecutive_count": self._consecutive_count,
                        "max_consecutive": self.max_consecutive,
                    },
                )
        return None

    def reset(self) -> None:
        """Reset consecutive tracking state."""
        super().reset()
        self._last_source = None
        self._consecutive_count = 0
