"""Token usage termination condition."""

from typing import Optional, Sequence

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

from .base import BaseTermination


class TokenUsageTermination(BaseTermination):
    """
    Terminates when cumulative token usage exceeds budget.

    Tracks total tokens (input + output) across all messages and
    stops when the budget is exceeded.

    Note: This requires messages to have usage information attached.
    In orchestration, this tracks tokens from all agent executions.

    Example:
        # Stop after 10000 tokens total
        termination = TokenUsageTermination(max_tokens=10000)

        # Combine with message limit
        termination = (
            TokenUsageTermination(5000) | MaxMessageTermination(20)
        )
    """

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

    def check(self, new_messages: Sequence[Message]) -> Optional[StopMessage]:
        """Check if token budget is exceeded."""
        # Count tokens from messages that have usage info
        for message in new_messages:
            if hasattr(message, "usage") and message.usage:
                usage = message.usage
                if hasattr(usage, "total_tokens"):
                    self.token_count += usage.total_tokens
                elif hasattr(usage, "tokens_input") and hasattr(usage, "tokens_output"):
                    self.token_count += usage.tokens_input + usage.tokens_output

        if self.token_count >= self.max_tokens:
            return self._set_termination(
                f"Token budget exceeded ({self.token_count}/{self.max_tokens})",
                {
                    "token_count": self.token_count,
                    "max_tokens": self.max_tokens,
                },
            )

        return None

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