"""Token budget resolution and status formatting.

Pure computation: given usage and limits, determine budget state.
No I/O imports.
"""

from deal import post, pre

from anima.core.schemas import JobSpec
from anima.core.tracker import BudgetTracker


# Note: Type checks (isinstance) are redundant with type annotations.
# job and persona are Pydantic models - structure validated at construction.
@pre(
    lambda cli_max_tokens, job: (
        (cli_max_tokens is None or cli_max_tokens >= 1)
        and (job.max_tokens is None or job.max_tokens >= 1)
    )
)
@post(lambda result: result is None or result >= 1)
def resolve_budget(cli_max_tokens: int | None, job: JobSpec) -> int | None:
    """Resolve token budget with precedence: CLI > job > None.

    Precedence order:
        1. cli_max_tokens (if provided, >= 1)
        2. job.max_tokens (if set in JobSpec)
        3. None (no budget constraint)

    Args:
        cli_max_tokens: CLI override for max tokens. Must be None or >= 1.
        job: JobSpec containing max_tokens override.

    Returns:
        Resolved token budget (>= 1) or None for unlimited.
        Returns None if all sources are unset or invalid.

    Doctests:
        >>> from anima.core.schemas import JobSpec

        # Precedence 1: CLI override takes priority
        >>> job = JobSpec(persona_ref="test@v1", instance_id="inst_1")
        >>> resolve_budget(1000, job)
        1000

        >>> # CLI override even when job has max_tokens
        >>> job_with_tokens = JobSpec(
        ...     persona_ref="test@v1", instance_id="inst_2", max_tokens=500
        ... )
        >>> resolve_budget(2000, job_with_tokens)
        2000

        # Precedence 2: Job max_tokens when no CLI override
        >>> resolve_budget(None, job_with_tokens)
        500

        # Precedence 3: None when no CLI or job override
        >>> resolve_budget(None, job) is None
        True

        # Contract enforcement: cli_max_tokens must be >= 1 if set
        >>> resolve_budget(0, job)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...

        # Contract enforcement: cli_max_tokens negative rejected
        >>> resolve_budget(-1, job)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...

        # Integration: full chain with all sources set
        >>> full_job = JobSpec(
        ...     persona_ref="full@v1",
        ...     instance_id="inst_full",
        ...     max_tokens=200
        ... )
        >>> resolve_budget(300, full_job)
        300
    """
    # Precedence 1: CLI override
    if cli_max_tokens is not None:
        return cli_max_tokens

    # Precedence 2: Job max_tokens
    if job.max_tokens is not None:
        return job.max_tokens

    # Precedence 3: No budget constraint
    return None


# =============================================================================
# Budget Message Formatting (Core - Pure Logic)
# =============================================================================


@pre(lambda grace_budget: grace_budget >= 0)
@post(lambda result: len(result) > 0)
@post(lambda result: "ctl_finish" in result)
def grace_system_message(grace_budget: int) -> str:
    """Generate the grace window system message.

    The message instructs the agent to use ctl_finish for graceful termination
    when the grace budget is exhausted.

    Args:
        grace_budget: The number of tokens in the grace window. Must be >= 0.

    Returns:
        A system message instructing about grace window usage.

    Doctests:
        >>> msg = grace_system_message(500)
        >>> "ctl_finish" in msg
        True
        >>> "grace" in msg.lower()
        True
        >>> "500" in msg
        True

        >>> # Zero grace budget still produces valid message
        >>> msg_zero = grace_system_message(0)
        >>> "ctl_finish" in msg_zero
        True
        >>> len(msg_zero) > 0
        True

        >>> # Contract: grace_budget must be non-negative
        >>> grace_system_message(-1)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...

        >>> # Message content verification
        >>> msg = grace_system_message(1000)
        >>> "grace window" in msg.lower()
        True
        >>> "token" in msg.lower()
        True
    """
    return (
        f"BUDGET GRACE WINDOW ACTIVE: You have entered the grace window with "
        f"{grace_budget} remaining tokens. "
        f"When the grace budget is exhausted, you MUST call ctl_finish to "
        f"gracefully terminate the session. Complete your current work and "
        f"use ctl_finish to end cleanly."
    )


@pre(lambda tracker: tracker is not None)
@post(lambda result: len(result) > 0)
@post(lambda result: "Budget:" in result or "budget:" in result.lower())
def format_budget_status(tracker: BudgetTracker) -> str:
    """Format budget status as a human-readable string.

    Handles three distinct states:
    - Normal: Budget is set, not in grace, shows remaining tokens
    - Unlimited: No budget constraint (budget_limit is None)
    - Grace: Currently in grace window, shows grace budget status

    Args:
        tracker: A BudgetTracker instance. Must not be None.

    Returns:
        A formatted status string describing the current budget state.

    Doctests:
        >>> # Normal budget: tokens remaining
        >>> tracker = BudgetTracker(1000)
        >>> tracker.record_usage(300)
        >>> status = format_budget_status(tracker)
        >>> "Budget:" in status
        True
        >>> "700" in status  # remaining
        True
        >>> "300" in status  # used
        True
        >>> "grace" in status.lower()
        False

        >>> # Normal budget: exhausted (not in grace yet)
        >>> exhausted = BudgetTracker(100)
        >>> exhausted.record_usage(100)
        >>> status = format_budget_status(exhausted)
        >>> "exhausted" in status.lower() or "0" in status
        True

        >>> # Unlimited budget
        >>> unlimited = BudgetTracker(None)
        >>> status = format_budget_status(unlimited)
        >>> "unlimited" in status.lower()
        True

        >>> unlimited.record_usage(10000)
        >>> status_unlimited = format_budget_status(unlimited)
        >>> "unlimited" in status_unlimited.lower()
        True

        >>> # Grace state: in grace window
        >>> grace_tracker = BudgetTracker(100, grace_budget=50)
        >>> grace_tracker.record_usage(100)
        >>> grace_tracker.enter_grace()
        >>> grace_tracker.record_usage(20)
        >>> status = format_budget_status(grace_tracker)
        >>> "grace" in status.lower()
        True
        >>> "30" in status  # 50 - 20 = 30 remaining in grace
        True

        >>> # Grace state: just entered grace
        >>> fresh_grace = BudgetTracker(200, grace_budget=100)
        >>> fresh_grace.record_usage(200)
        >>> fresh_grace.enter_grace()
        >>> status = format_budget_status(fresh_grace)
        >>> "grace" in status.lower()
        True
        >>> "100" in status  # full grace budget available
        True

        >>> # Contract: tracker must not be None
        >>> format_budget_status(None)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...

        >>> # Grace exhausted (but still in grace state)
        >>> grace_done = BudgetTracker(100, grace_budget=30)
        >>> grace_done.record_usage(100)
        >>> grace_done.enter_grace()
        >>> grace_done.record_usage(30)
        >>> status = format_budget_status(grace_done)
        >>> "grace" in status.lower()
        True
    """
    # Unlimited budget
    if tracker.budget_limit is None:
        return "Budget: unlimited"

    # Grace state
    if tracker.in_grace:
        grace_used = tracker.used - tracker.grace_start_usage
        grace_remaining = tracker.grace_budget - grace_used
        return (
            f"Budget: GRACE MODE - "
            f"{grace_used}/{tracker.grace_budget} grace tokens used, "
            f"{max(0, grace_remaining)} remaining"
        )

    # Normal budget
    remaining = tracker.remaining()
    if remaining is not None and remaining > 0:
        return (
            f"Budget: {tracker.used}/{tracker.budget_limit} tokens used, "
            f"{remaining} remaining"
        )
    else:
        return f"Budget: {tracker.used}/{tracker.budget_limit} tokens used, exhausted"


# Re-export for backward compatibility
__all__ = [
    "resolve_budget",
    "grace_system_message",
    "format_budget_status",
    "BudgetTracker",
]
