"""
Rollout Compression - manages message history eviction for infinite-running agents.

As agents run, their conversation history grows. These compressors implement
different strategies for evicting old messages to keep context bounded.

Available compressors:
- NoOpRolloutCompressor: No eviction (context will eventually overflow)
- LRURolloutCompressor: Evict oldest messages, keeping system prompt
- SlidingWindowRolloutCompressor: Keep last N turns (user+assistant pairs)
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any


@dataclass
class EvictionStats:
    """Statistics about message eviction."""

    messages_before: int = 0
    messages_after: int = 0
    messages_evicted: int = 0
    tokens_evicted_estimate: int = 0  # Rough estimate
    eviction_triggered: bool = False

    def to_dict(self) -> dict[str, Any]:
        return {
            "messages_before": self.messages_before,
            "messages_after": self.messages_after,
            "messages_evicted": self.messages_evicted,
            "tokens_evicted_estimate": self.tokens_evicted_estimate,
            "eviction_triggered": self.eviction_triggered,
        }


class RolloutCompressor(ABC):
    """
    Abstract base class for rollout (message history) compression.

    Subclasses implement different eviction strategies to keep the
    conversation history bounded while preserving important context.
    """

    name: str = "base"

    @abstractmethod
    def compress(self, messages: list[dict]) -> tuple[list[dict], EvictionStats]:
        """
        Compress message history, returning trimmed version and stats.

        Args:
            messages: List of message dicts with 'role' and 'content' keys

        Returns:
            (compressed_messages, eviction_stats)
        """
        pass

    def should_compress(self, messages: list[dict]) -> bool:
        """Check if compression is needed (subclasses may override)."""
        return True

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"


class NoOpRolloutCompressor(RolloutCompressor):
    """
    No-op compressor - keeps all messages.

    Use this when you want to disable rollout compression and let the
    context window naturally overflow (will error eventually).
    """

    name = "noop"

    def compress(self, messages: list[dict]) -> tuple[list[dict], EvictionStats]:
        """Pass through unchanged."""
        return messages, EvictionStats(
            messages_before=len(messages),
            messages_after=len(messages),
            eviction_triggered=False,
        )


class LRURolloutCompressor(RolloutCompressor):
    """
    LRU (Least Recently Used) compressor - evicts oldest messages.

    Keeps the system prompt and the most recent messages. When the message
    count exceeds max_messages, evicts oldest non-system messages.

    Args:
        max_messages: Maximum messages to keep (default: 50)
        preserve_system: Keep all system messages (default: True)
        preserve_first_user: Keep first user message as context (default: True)
    """

    name = "lru"

    def __init__(
        self,
        max_messages: int = 50,
        preserve_system: bool = True,
        preserve_first_user: bool = True,
    ):
        self.max_messages = max_messages
        self.preserve_system = preserve_system
        self.preserve_first_user = preserve_first_user

    def should_compress(self, messages: list[dict]) -> bool:
        """Only compress if we exceed max_messages."""
        return len(messages) > self.max_messages

    def compress(self, messages: list[dict]) -> tuple[list[dict], EvictionStats]:
        """Evict oldest messages, keeping system prompt and recent history."""
        stats = EvictionStats(messages_before=len(messages))

        if not self.should_compress(messages):
            stats.messages_after = len(messages)
            return messages, stats

        # Separate preserved messages from evictable ones
        preserved = []
        evictable = []

        first_user_seen = False
        for i, msg in enumerate(messages):
            role = msg.get("role", "")

            # Always preserve system messages
            if self.preserve_system and role == "system":
                preserved.append((i, msg))
            # Preserve first user message as task context
            elif self.preserve_first_user and role == "user" and not first_user_seen:
                preserved.append((i, msg))
                first_user_seen = True
            else:
                evictable.append((i, msg))

        # Calculate how many evictable messages to keep
        preserved_count = len(preserved)
        keep_count = max(0, self.max_messages - preserved_count)

        # Keep the most recent evictable messages
        kept_evictable = evictable[-keep_count:] if keep_count > 0 else []
        evicted = evictable[:-keep_count] if keep_count > 0 and len(evictable) > keep_count else []

        # Merge preserved and kept messages, maintaining original order
        all_kept = preserved + kept_evictable
        all_kept.sort(key=lambda x: x[0])  # Sort by original index
        result = [msg for _, msg in all_kept]

        # Estimate tokens evicted (rough: ~4 chars per token)
        evicted_content = sum(len(str(msg.get("content", ""))) for _, msg in evicted)
        tokens_evicted = evicted_content // 4

        stats.messages_after = len(result)
        stats.messages_evicted = len(evicted)
        stats.tokens_evicted_estimate = tokens_evicted
        stats.eviction_triggered = len(evicted) > 0

        return result, stats

    def __repr__(self) -> str:
        return f"LRURolloutCompressor(max_messages={self.max_messages})"


class SlidingWindowRolloutCompressor(RolloutCompressor):
    """
    Sliding window compressor - keeps last N turns (user+assistant pairs).

    A "turn" is a user message followed by an assistant response. This
    preserves conversation coherence better than raw message count.

    Args:
        max_turns: Maximum turns to keep (default: 20)
        preserve_system: Keep all system messages (default: True)
        preserve_first_turn: Keep first turn as context (default: True)
    """

    name = "sliding_window"

    def __init__(
        self,
        max_turns: int = 20,
        preserve_system: bool = True,
        preserve_first_turn: bool = True,
    ):
        self.max_turns = max_turns
        self.preserve_system = preserve_system
        self.preserve_first_turn = preserve_first_turn

    def compress(self, messages: list[dict]) -> tuple[list[dict], EvictionStats]:
        """Keep last N turns, preserving system messages."""
        stats = EvictionStats(messages_before=len(messages))

        # Extract system messages
        system_messages = []
        conversation = []

        for msg in messages:
            if msg.get("role") == "system":
                system_messages.append(msg)
            else:
                conversation.append(msg)

        # Group conversation into turns (user + assistant + tool results)
        turns: list[list[dict]] = []
        current_turn: list[dict] = []

        for msg in conversation:
            role = msg.get("role", "")
            if role == "user" and current_turn:
                # New user message starts a new turn
                turns.append(current_turn)
                current_turn = [msg]
            else:
                current_turn.append(msg)

        # Don't forget the last turn
        if current_turn:
            turns.append(current_turn)

        # Decide which turns to keep
        if len(turns) <= self.max_turns:
            # No eviction needed
            result = system_messages + conversation
            stats.messages_after = len(result)
            return result, stats

        # Keep first turn + last (max_turns - 1) turns
        kept_turns = []
        if self.preserve_first_turn and turns:
            kept_turns.append(turns[0])
            remaining_turns = turns[1:]
            kept_turns.extend(remaining_turns[-(self.max_turns - 1):])
        else:
            kept_turns = turns[-self.max_turns:]

        # Flatten kept turns back into messages
        kept_conversation = []
        for turn in kept_turns:
            kept_conversation.extend(turn)

        result = system_messages + kept_conversation

        # Calculate eviction stats
        evicted_count = len(messages) - len(result)
        stats.messages_after = len(result)
        stats.messages_evicted = evicted_count
        stats.eviction_triggered = evicted_count > 0

        return result, stats

    def __repr__(self) -> str:
        return f"SlidingWindowRolloutCompressor(max_turns={self.max_turns})"


# =============================================================================
# Factory
# =============================================================================


def get_rollout_compressor(
    name: str = "lru",
    **kwargs,
) -> RolloutCompressor:
    """
    Get a rollout compressor by name.

    Args:
        name: Compressor name ("noop", "lru", "sliding_window")
        **kwargs: Passed to compressor constructor

    Returns:
        Configured RolloutCompressor instance
    """
    compressors = {
        "noop": NoOpRolloutCompressor,
        "lru": LRURolloutCompressor,
        "sliding_window": SlidingWindowRolloutCompressor,
    }

    if name not in compressors:
        raise ValueError(f"Unknown rollout compressor: {name}. Available: {list(compressors.keys())}")

    return compressors[name](**kwargs)
