"""Context compression strategies."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any

from thryve.core.context.models import Message
from thryve.utils import estimate_tokens, get_logger

if TYPE_CHECKING:
    from thryve.core.backends.adapter import BackendAdapter

logger = get_logger("context.compression")


class CompressionStrategy(ABC):
    """Base class for context compression strategies."""

    @abstractmethod
    def compress(self, messages: list[Message], target_tokens: int) -> list[Message]:
        """Return a compressed copy of *messages* that fits within *target_tokens*."""


class TruncateStrategy(CompressionStrategy):
    """Drop the oldest non-system messages until the budget is met.

    System messages are always retained.
    """

    def compress(self, messages: list[Message], target_tokens: int) -> list[Message]:
        # Split system messages from the rest.
        system: list[Message] = [m for m in messages if m.role == "system"]
        others: list[Message] = [m for m in messages if m.role != "system"]

        total = sum(estimate_tokens(m.text) for m in system)
        keep: list[Message] = []

        # Walk from newest to oldest, adding messages until budget exhausted.
        for msg in reversed(others):
            cost = estimate_tokens(msg.text)
            if total + cost > target_tokens:
                break
            keep.append(msg)
            total += cost

        keep.reverse()
        return system + keep


class SummarizeStrategy(CompressionStrategy):
    """Summarize older messages via the LLM, keeping recent ones verbatim.

    Requires an LLM backend to generate the summary.
    """

    def __init__(self, llm: BackendAdapter, recent_keep: int = 4) -> None:
        self._llm = llm
        self._recent_keep = recent_keep

    def compress(self, messages: list[Message], target_tokens: int) -> list[Message]:
        if len(messages) <= self._recent_keep + 1:
            return list(messages)

        system = [m for m in messages if m.role == "system"]
        non_system = [m for m in messages if m.role != "system"]

        # Split into "to summarize" and "to keep verbatim".
        to_summarize = non_system[: -self._recent_keep]
        to_keep = non_system[-self._recent_keep :]

        # Build a summary prompt
        text_block = "\n".join(f"[{m.role}]: {m.text}" for m in to_summarize)
        summary_prompt = (
            "Summarize the following conversation concisely, preserving key "
            "facts, decisions, and tool results:\n\n" + text_block
        )

        try:
            response = self._llm.generate(
                [Message(role="user", content=summary_prompt)]
            )
            summary_msg = Message(role="system", content=f"[Conversation summary] {response.content}")
        except Exception:
            logger.warning("Summarize failed; falling back to truncation")
            return TruncateStrategy().compress(messages, target_tokens)

        return system + [summary_msg] + to_keep


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

def get_strategy(name: str, **kwargs: Any) -> CompressionStrategy:
    """Return a compression strategy by *name*."""
    if name == "truncate":
        return TruncateStrategy()
    elif name == "summarize":
        llm = kwargs.get("llm")
        if llm is None:
            raise ValueError("SummarizeStrategy requires an 'llm' keyword argument")
        return SummarizeStrategy(llm, recent_keep=kwargs.get("recent_keep", 4))
    else:
        raise ValueError(f"Unknown compression strategy: {name!r}")
