"""cortex/types/budget.py — Budget allocation and compression types."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict

from cortex.types.enums import CompressionLevel


@dataclass
class BudgetAllocation:
    """Result of a budget allocation pass.

    Records how many tokens were allocated per layer and what was truncated.
    """

    total_budget: int
    """Total token budget available."""

    cornerstone_reserved: int
    """Tokens reserved exclusively for cornerstone memories."""

    allocations: Dict[str, int] = field(default_factory=dict)
    """Tokens allocated per layer key (MemoryLayer.value)."""

    truncated: Dict[str, int] = field(default_factory=dict)
    """Tokens that were truncated per layer key due to budget exhaustion."""

    @property
    def used(self) -> int:
        """Total tokens consumed (reserved + allocated)."""
        return self.cornerstone_reserved + sum(self.allocations.values())

    @property
    def remaining(self) -> int:
        """Tokens still available for allocation."""
        return max(0, self.total_budget - self.used)

    @property
    def utilization(self) -> float:
        """Fraction of budget consumed (0.0–1.0+)."""
        if self.total_budget == 0:
            return 0.0
        return self.used / self.total_budget

    def is_over_budget(self) -> bool:
        """True if allocations exceeded the total budget."""
        return self.used > self.total_budget


@dataclass
class CompressedItem:
    """A memory item at a specific compression level, with token accounting.

    Produced by TokenBudgetManager.compress_to_fit() and consumed by the
    injection layer (CortexPlugin.assemble_context).
    """

    content: str
    compression: CompressionLevel
    token_count: int
    original_tokens: int

    @property
    def compression_ratio(self) -> float:
        """Fraction of original tokens retained (0.0 = nothing, 1.0 = full)."""
        if self.original_tokens == 0:
            return 1.0
        return self.token_count / self.original_tokens

    @property
    def tokens_saved(self) -> int:
        """Tokens saved relative to original content."""
        return max(0, self.original_tokens - self.token_count)
