"""Short-term memory management (unbounded deque, compression driven by token limits)."""

from collections import deque
from typing import List

from llm.base import LLMMessage


class ShortTermMemory:
    """Manages recent messages in an unbounded deque."""

    def __init__(self):
        """Initialize short-term memory (unbounded; compression is driven by token limits)."""
        self.messages: deque[LLMMessage] = deque()

    def add_message(self, message: LLMMessage) -> None:
        """Add a message to short-term memory.

        Args:
            message: LLMMessage to add
        """
        self.messages.append(message)

    def get_messages(self) -> List[LLMMessage]:
        """Get all messages in short-term memory.

        Returns:
            List of messages, oldest to newest
        """
        return list(self.messages)

    def clear(self) -> List[LLMMessage]:
        """Clear all messages and return them.

        Returns:
            List of all messages that were cleared
        """
        messages = list(self.messages)
        self.messages.clear()
        return messages

    def remove_first(self, count: int) -> List[LLMMessage]:
        """Remove the first N messages (oldest) from memory.

        This is useful after compression to remove only the compressed messages
        while preserving any new messages that arrived during compression.

        Args:
            count: Number of messages to remove from the front

        Returns:
            List of removed messages
        """
        return [self.messages.popleft() for _ in range(min(count, len(self.messages)))]

    def count(self) -> int:
        """Get current message count.

        Returns:
            Number of messages in short-term memory
        """
        return len(self.messages)

    def remove_last(self, count: int = 1) -> None:
        """Remove the last N messages (newest) from memory.

        This is useful for rolling back incomplete exchanges (e.g., after interruption).

        Args:
            count: Number of messages to remove from the end (default: 1)
        """
        for _ in range(min(count, len(self.messages))):
            self.messages.pop()
