"""Budget tracker with grace window support.

Pure computation for token budget tracking.
No I/O imports.
"""

from deal import post, pre


class BudgetTracker:
    """Token budget tracker with grace window support.

    Tracks token usage against a budget limit, with an optional grace window
    that provides additional tokens after exhaustion.

    Attributes:
        budget_limit: Maximum tokens allowed, or None for unlimited.
        grace_budget: Extra tokens available in grace window.
        used: Total tokens consumed so far.
        in_grace: Whether currently in grace window.
        grace_start_usage: Token count when grace was entered.
    """

    __slots__ = (
        "budget_limit",
        "grace_budget",
        "used",
        "in_grace",
        "grace_start_usage",
    )

    @pre(
        lambda self, budget_limit, grace_budget=500: (
            (budget_limit is None or budget_limit >= 1) and grace_budget >= 0
        )
    )
    def __init__(self, budget_limit: int | None, grace_budget: int = 500) -> None:
        """Initialize budget tracker.

        Args:
            budget_limit: Maximum tokens allowed, or None for unlimited budget.
            grace_budget: Extra tokens available after exhaustion. Defaults to 500.

        Doctests:
            >>> # Normal budget
            >>> tracker = BudgetTracker(1000)
            >>> tracker.budget_limit
            1000
            >>> tracker.grace_budget
            500
            >>> tracker.used
            0
            >>> tracker.in_grace
            False

            >>> # Unlimited budget
            >>> unlimited = BudgetTracker(None)
            >>> unlimited.budget_limit is None
            True

            >>> # Custom grace budget
            >>> custom = BudgetTracker(100, grace_budget=200)
            >>> custom.grace_budget
            200

            >>> # Zero grace budget (no grace allowed)
            >>> no_grace = BudgetTracker(100, grace_budget=0)
            >>> no_grace.grace_budget
            0

            >>> # Contract: budget_limit must be >= 1 if not None
            >>> BudgetTracker(0)  # doctest: +IGNORE_EXCEPTION_DETAIL
            Traceback (most recent call last):
            ...
            PreContractError: ...

            >>> # Contract: grace_budget must be >= 0
            >>> BudgetTracker(100, grace_budget=-1)  # doctest: +IGNORE_EXCEPTION_DETAIL
            Traceback (most recent call last):
            ...
            PreContractError: ...
        """
        self.budget_limit = budget_limit
        self.grace_budget = grace_budget
        self.used = 0
        self.in_grace = False
        self.grace_start_usage = 0

    @pre(lambda self, tokens: tokens >= 0)
    @post(lambda result: result is None)
    def record_usage(self, tokens: int) -> None:
        """Record token usage.

        Args:
            tokens: Number of tokens consumed. Must be >= 0.

        Doctests:
            >>> tracker = BudgetTracker(1000)
            >>> tracker.record_usage(100)
            >>> tracker.used
            100

            >>> # Multiple recordings accumulate
            >>> tracker.record_usage(50)
            >>> tracker.used
            150

            >>> # Zero tokens is valid (no-op)
            >>> tracker.record_usage(0)
            >>> tracker.used
            150

            >>> # Contract: tokens must be >= 0
            >>> tracker.record_usage(-1)  # doctest: +IGNORE_EXCEPTION_DETAIL
            Traceback (most recent call last):
            ...
            PreContractError: ...
        """
        self.used += tokens

    @post(lambda result: result is None or result >= 0)
    def remaining(self) -> int | None:
        """Get remaining tokens before exhaustion.

        Returns:
            Remaining tokens, or None if budget is unlimited.
            In normal mode: budget_limit - used.
            In grace mode: grace_start_usage + grace_budget - used.
            If past limit: returns 0 (never negative).

        Doctests:
            >>> # Normal budget
            >>> tracker = BudgetTracker(1000)
            >>> tracker.remaining()
            1000

            >>> tracker.record_usage(300)
            >>> tracker.remaining()
            700

            >>> # Unlimited budget always returns None
            >>> unlimited = BudgetTracker(None)
            >>> unlimited.remaining() is None
            True

            >>> unlimited.record_usage(1000)
            >>> unlimited.remaining() is None
            True

            >>> # Grace window: remaining resets to grace_budget
            >>> tracker = BudgetTracker(100, grace_budget=50)
            >>> tracker.record_usage(100)  # Exhaust normal budget
            >>> tracker.enter_grace()
            >>> tracker.remaining()  # Grace budget available
            50

            >>> tracker.record_usage(30)
            >>> tracker.remaining()  # Grace remaining
            20

            >>> # Past limit returns 0
            >>> tracker.record_usage(25)
            >>> tracker.remaining()
            0
        """
        # Unlimited budget
        if self.budget_limit is None:
            return None

        # In grace window
        if self.in_grace:
            remaining = self.grace_start_usage + self.grace_budget - self.used
            return max(0, remaining)

        # Normal budget
        remaining = self.budget_limit - self.used
        return max(0, remaining)

    @post(lambda result: isinstance(result, bool))
    def is_exhausted(self) -> bool:
        """Check if budget is exhausted.

        In grace mode: checks if grace budget is exhausted.
        In normal mode: checks if normal budget is exhausted.

        Returns:
            True if no tokens remain (in current mode).

        Doctests:
            >>> # Normal budget
            >>> tracker = BudgetTracker(100)
            >>> tracker.is_exhausted()
            False

            >>> tracker.record_usage(100)
            >>> tracker.is_exhausted()
            True

            >>> # Unlimited never exhausted
            >>> unlimited = BudgetTracker(None)
            >>> unlimited.is_exhausted()
            False

            >>> unlimited.record_usage(1000000)
            >>> unlimited.is_exhausted()
            False

            >>> # Grace window: exhaustion resets
            >>> tracker = BudgetTracker(100, grace_budget=50)
            >>> tracker.record_usage(100)
            >>> tracker.is_exhausted()  # Normal exhausted
            True

            >>> tracker.enter_grace()
            >>> tracker.is_exhausted()  # Grace gives more
            False

            >>> tracker.record_usage(50)  # Use all grace
            >>> tracker.is_exhausted()  # Grace exhausted
            True
        """
        # Unlimited never exhausted
        if self.budget_limit is None:
            return False

        # In grace window
        if self.in_grace:
            return self.used >= self.grace_start_usage + self.grace_budget

        # Normal budget
        return self.used >= self.budget_limit

    @post(lambda result: isinstance(result, bool))
    def should_enter_grace(self) -> bool:
        """Check if we should enter grace window.

        Grace should be entered when:
        - Not already in grace
        - Normal budget is exhausted
        - Grace budget is available (> 0)

        Returns:
            True if conditions are met to enter grace.

        Doctests:
            >>> # Normal budget not exhausted
            >>> tracker = BudgetTracker(100)
            >>> tracker.should_enter_grace()
            False

            >>> # Normal budget exhausted
            >>> tracker.record_usage(100)
            >>> tracker.should_enter_grace()
            True

            >>> # Already in grace - should not enter again
            >>> tracker.enter_grace()
            >>> tracker.should_enter_grace()
            False

            >>> # No grace budget available
            >>> no_grace = BudgetTracker(100, grace_budget=0)
            >>> no_grace.record_usage(100)
            >>> no_grace.should_enter_grace()
            False

            >>> # Unlimited never needs grace
            >>> unlimited = BudgetTracker(None)
            >>> unlimited.should_enter_grace()
            False
        """
        # Unlimited never needs grace
        if self.budget_limit is None:
            return False

        # Already in grace
        if self.in_grace:
            return False

        # No grace budget available
        if self.grace_budget <= 0:
            return False

        # Normal budget exhausted and grace available
        return self.used >= self.budget_limit

    @post(lambda result: result is None)
    def enter_grace(self) -> None:
        """Enter grace window if conditions are met.

        Idempotent: safe to call multiple times.
        Only enters grace if should_enter_grace() would return True.

        Sets grace_start_usage to current usage, giving grace_budget
        additional tokens from this point.

        Doctests:
            >>> # Enter grace when exhausted
            >>> tracker = BudgetTracker(100, grace_budget=50)
            >>> tracker.record_usage(100)
            >>> tracker.enter_grace()
            >>> tracker.in_grace
            True
            >>> tracker.grace_start_usage
            100

            >>> # Idempotent: calling again is safe
            >>> tracker.enter_grace()
            >>> tracker.in_grace
            True

            >>> # Cannot enter if not exhausted
            >>> tracker2 = BudgetTracker(100)
            >>> tracker2.record_usage(50)
            >>> tracker2.enter_grace()
            >>> tracker2.in_grace
            False

            >>> # Cannot enter with no grace budget
            >>> tracker3 = BudgetTracker(100, grace_budget=0)
            >>> tracker3.record_usage(100)
            >>> tracker3.enter_grace()
            >>> tracker3.in_grace
            False
        """
        if self.should_enter_grace():
            self.in_grace = True
            self.grace_start_usage = self.used

    @post(lambda result: isinstance(result, bool))
    def is_grace_exhausted(self) -> bool:
        """Check if grace budget is exhausted.

        Returns:
            True if in grace mode and all grace tokens used.
            False if not in grace or grace tokens remain.

        Doctests:
            >>> # Not in grace
            >>> tracker = BudgetTracker(100)
            >>> tracker.is_grace_exhausted()
            False

            >>> # In grace with tokens remaining
            >>> tracker.record_usage(100)
            >>> tracker.enter_grace()
            >>> tracker.is_grace_exhausted()
            False

            >>> # Grace exhausted
            >>> tracker.record_usage(500)  # Use all grace
            >>> tracker.is_grace_exhausted()
            True

            >>> # Not in grace even if normal exhausted
            >>> tracker2 = BudgetTracker(100)
            >>> tracker2.record_usage(100)
            >>> tracker2.is_grace_exhausted()  # Not in grace yet
            False

            >>> # Unlimited never in grace
            >>> unlimited = BudgetTracker(None)
            >>> unlimited.is_grace_exhausted()
            False
        """
        # Not in grace means grace cannot be exhausted
        if not self.in_grace:
            return False

        # Check if used exceeds grace budget
        return self.used >= self.grace_start_usage + self.grace_budget
