from __future__ import annotations

"""Token-aware memory backend."""

from dataclasses import dataclass, field
from math import ceil
from typing import Protocol

from agentrspy.core import Memory, Message, Role


class Tokenizer(Protocol):
    """Counts approximate tokens for a string."""

    def count(self, text: str) -> int:
        """Return the estimated token count."""


class ApproximateTokenizer:
    """Lightweight tokenizer approximation that avoids external dependencies."""

    def count(self, text: str) -> int:
        return ceil(len(text) / 3.5)


@dataclass(slots=True)
class TokenAwareMemory(Memory):
    """Memory backend that trims history to fit a token budget."""

    max_tokens: int
    tokenizer: Tokenizer = field(default_factory=ApproximateTokenizer)
    messages: list[Message] = field(default_factory=list)

    @classmethod
    def new(cls, max_tokens: int) -> "TokenAwareMemory":
        return cls(max_tokens=max_tokens)

    @classmethod
    def with_tokenizer(
        cls, max_tokens: int, tokenizer: Tokenizer
    ) -> "TokenAwareMemory":
        return cls(max_tokens=max_tokens, tokenizer=tokenizer)

    def _total_tokens(self) -> int:
        return sum(
            self.tokenizer.count(message.text_content()) for message in self.messages
        )

    def _trim_to_budget(self) -> None:
        while self._total_tokens() > self.max_tokens and len(self.messages) > 1:
            for index, message in enumerate(self.messages):
                if message.role is not Role.SYSTEM:
                    self.messages.pop(index)
                    break
            else:
                break

    async def store(self, key: str, value: Message) -> None:
        del key
        self.messages.append(value)
        self._trim_to_budget()

    async def retrieve(self, query: str, limit: int) -> list[Message]:
        lowered = query.lower()
        return [
            message
            for message in self.messages
            if lowered in message.text_content().lower()
        ][:limit]

    async def history(self) -> list[Message]:
        return list(self.messages)

    async def clear(self) -> None:
        self.messages.clear()

    async def token_count(self) -> int:
        return self._total_tokens()
