"""Context manager – tracks conversation messages, tokens, checkpoints."""

from __future__ import annotations

import uuid
from typing import Any

from thryve.core.context.checkpoint import Checkpoint
from thryve.core.context.compression import CompressionStrategy, TruncateStrategy
from thryve.core.context.models import Message, MessageRole
from thryve.utils import estimate_tokens, get_logger

logger = get_logger("context.manager")


class ContextManager:
    """Manages conversation context: messages, token budgets, and checkpoints.

    When the accumulated token count exceeds *max_tokens*, the manager
    automatically compacts older messages using the configured compression
    strategy.
    """

    def __init__(
        self,
        max_tokens: int = 8000,
        strategy: CompressionStrategy | None = None,
    ) -> None:
        self.messages: list[Message] = []
        self.max_tokens = max_tokens
        self._strategy: CompressionStrategy = strategy or TruncateStrategy()
        self._checkpoints: dict[str, Checkpoint] = {}

    # ------------------------------------------------------------------
    # Message management
    # ------------------------------------------------------------------

    def add_message(
        self,
        role: MessageRole,
        content: str,
        **kwargs: Any,
    ) -> Message:
        """Create and append a :class:`Message`."""
        msg = Message(role=role, content=content, **kwargs)
        msg.token_count = estimate_tokens(msg.text)
        self.messages.append(msg)
        self._check_and_compact()
        return msg

    def add(self, message: Message) -> None:
        """Append an existing :class:`Message` instance."""
        if message.token_count == 0:
            message.token_count = estimate_tokens(message.text)
        self.messages.append(message)
        self._check_and_compact()

    @property
    def total_tokens(self) -> int:
        return sum(m.token_count for m in self.messages)

    # ------------------------------------------------------------------
    # Compaction
    # ------------------------------------------------------------------

    def _check_and_compact(self) -> None:
        if self.total_tokens > self.max_tokens:
            logger.info(
                "Token budget exceeded (%d > %d) – compacting",
                self.total_tokens,
                self.max_tokens,
            )
            self.compact()

    def compact(self) -> None:
        """Compress messages using the configured strategy."""
        self.messages = self._strategy.compress(self.messages, self.max_tokens)
        # Re-estimate tokens for the resulting messages.
        for m in self.messages:
            m.token_count = estimate_tokens(m.text)

    # ------------------------------------------------------------------
    # Checkpoints
    # ------------------------------------------------------------------

    def create_checkpoint(self, name: str = "") -> str:
        """Snapshot the current messages; return the checkpoint id."""
        cid = uuid.uuid4().hex
        self._checkpoints[cid] = Checkpoint(
            id=cid,
            name=name,
            messages=[m for m in self.messages],  # shallow copy of list
        )
        return cid

    def rollback_to_checkpoint(self, checkpoint_id: str) -> None:
        """Restore messages from a checkpoint."""
        cp = self._checkpoints.get(checkpoint_id)
        if cp is None:
            raise KeyError(f"Checkpoint {checkpoint_id!r} not found")
        self.messages = list(cp.messages)

    def list_checkpoints(self) -> list[Checkpoint]:
        return list(self._checkpoints.values())

    # ------------------------------------------------------------------
    # Misc
    # ------------------------------------------------------------------

    def clear(self) -> None:
        """Remove all messages (but keep checkpoints)."""
        self.messages.clear()

    def get_messages(self, last_n: int | None = None) -> list[Message]:
        """Return messages, optionally the last *n* only."""
        if last_n is not None:
            return self.messages[-last_n:]
        return list(self.messages)
