"""Priority-based pruning strategy."""

from __future__ import annotations

from typing import TYPE_CHECKING

from context_manager.strategies.base import PruningStrategy

if TYPE_CHECKING:
    from context_manager.models import ContextBlock


class PriorityPruning(PruningStrategy):
    """Drop lowest-priority blocks first until the budget is satisfied.

    Blocks with the same priority are dropped in reverse insertion order
    (most recently added blocks are dropped first).
    """

    def prune(
        self,
        blocks: list[ContextBlock],
        token_budget: int,
    ) -> list[ContextBlock]:
        """Return blocks fitting within *token_budget*, highest priority first."""
        if not blocks:
            return []

        # Sort: highest priority first, then original order (stable sort).
        sorted_blocks = sorted(blocks, key=lambda b: b.priority, reverse=True)

        kept: list[ContextBlock] = []
        used = 0

        for block in sorted_blocks:
            cost = block.token_count or 0
            if used + cost <= token_budget:
                kept.append(block)
                used += cost

        # Restore original insertion order for the final prompt.
        original_order = {b.id: idx for idx, b in enumerate(blocks)}
        kept.sort(key=lambda b: original_order[b.id])

        return kept
