"""Budget and token tracking helpers for AgentRunner.

Extracted from agent.py to reduce file size.
Spec source: vectl step ``shrink-agent-py`` in ``gate-review-debt`` phase.
"""

# @invar:allow shell_result: Budget helpers return Result types
# @shell_orchestration: Budget helpers delegate to state I/O for logging and transitions

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from returns.result import Failure, Result

from anima.core.lifecycle import LifecycleEvent
from anima.core.schemas import (
    AgentResult,
    AnimaError,
    BudgetExhaustedError,
    MaxTurnsReachedError,
    TokenUsage,
)
from anima.shell.state import append_log, transition_state, update_agent_state

if TYPE_CHECKING:
    from anima.core.tracker import BudgetTracker


# @shell_orchestration: Delegates to update_agent_state I/O
def update_state(instance_id: str, turn_count: int, token_usage: TokenUsage) -> None:
    """Update state file with current turn count and token usage.

    Args:
        instance_id: The agent's unique identifier.
        turn_count: Current turn number.
        token_usage: Cumulative token usage.
    """
    update_agent_state(
        instance_id,
        {
            "current_turn": turn_count,
            "token_usage": token_usage.model_dump(),
        },
    )


# @shell_orchestration: Delegates to budget_tracker and update_agent_state I/O
def track_token_usage(
    usage: Any,
    token_usage: TokenUsage,
    budget_tracker: BudgetTracker,
    instance_id: str,
    turn_count: int,
) -> None:
    """Track token usage from pydantic-ai run.

    Mutates ``token_usage`` and ``budget_tracker`` in place, then persists
    updated state.

    Args:
        usage: Usage object from pydantic-ai run with input_tokens/output_tokens.
        token_usage: Mutable TokenUsage accumulator on the runner.
        budget_tracker: Budget tracker instance on the runner.
        instance_id: The agent's unique identifier.
        turn_count: Current turn number.
    """
    if usage:
        prompt_tokens = usage.input_tokens or 0
        completion_tokens = usage.output_tokens or 0

        # Update token usage tracking
        token_usage.prompt_tokens += prompt_tokens
        token_usage.completion_tokens += completion_tokens
        # Track the most recent request's prompt size for compaction decisions.
        # Cumulative prompt_tokens is the sum across all requests in the session;
        # last_request_prompt_tokens is the actual input size of the latest request,
        # which is what determines if we're near the context window limit.
        token_usage.last_request_prompt_tokens = prompt_tokens

        # Track in budget tracker
        total_tokens = prompt_tokens + completion_tokens
        budget_tracker.record_usage(total_tokens)

        # Update state
        update_state(instance_id, turn_count, token_usage)


# @shell_complexity: Budget threshold checks with grace period logic
def check_budget_exhaustion(
    budget_tracker: BudgetTracker,
    in_grace_window: bool,
    instance_id: str,
) -> tuple[Result[AgentResult, AnimaError] | None, bool]:
    """Check for budget exhaustion and grace window entry.

    Args:
        budget_tracker: Budget tracker instance.
        in_grace_window: Whether the agent is currently in grace window.
        instance_id: The agent's unique identifier.

    Returns:
        Tuple of (result_if_terminal, new_in_grace_window).
        result_if_terminal is None when the loop should continue.
        new_in_grace_window reflects whether grace was entered.
    """
    # Unlimited budget never triggers
    if budget_tracker.budget_limit is None:
        return None, in_grace_window

    # Already in grace window
    if in_grace_window:
        if budget_tracker.is_grace_exhausted():
            # Grace exhausted without ctl_finish - ERROR
            transition_state(
                instance_id,
                LifecycleEvent.budget_exhausted,
                error="Grace budget exhausted without ctl_finish",
            )
            append_log(
                instance_id,
                "grace_exhausted",
                {
                    "budget_used": budget_tracker.used,
                    "grace_budget": budget_tracker.grace_budget,
                },
            )
            return (
                Failure(
                    BudgetExhaustedError(
                        "Grace budget exhausted without ctl_finish",
                        details={
                            "budget_used": budget_tracker.used,
                            "grace_budget": budget_tracker.grace_budget,
                        },
                    )
                ),
                in_grace_window,
            )
        return None, in_grace_window

    # Check if should enter grace
    if budget_tracker.should_enter_grace():
        budget_tracker.enter_grace()

        # Log grace entry
        append_log(
            instance_id,
            "grace_entered",
            {
                "budget_used": budget_tracker.used,
                "grace_budget": budget_tracker.grace_budget,
            },
        )

        # Update state with grace mode metadata (no state change)
        update_agent_state(instance_id, {})

        return None, True

    # Check if exhausted (no grace available)
    if budget_tracker.is_exhausted():
        transition_state(
            instance_id,
            LifecycleEvent.budget_exhausted,
            error=f"Budget exhausted: {budget_tracker.used} tokens",
        )
        return (
            Failure(
                BudgetExhaustedError(
                    f"Budget exhausted: {budget_tracker.used} tokens",
                    details={
                        "budget_used": budget_tracker.used,
                        "budget_limit": budget_tracker.budget_limit,
                    },
                )
            ),
            in_grace_window,
        )

    return None, in_grace_window


def handle_model_request_turn(
    turn_count: int,
    max_turns: int | None,
    in_grace_window: bool,
    budget_tracker: BudgetTracker,
    instance_id: str,
) -> tuple[Result[AgentResult, AnimaError] | None, bool, bool]:
    """Handle a model request turn. Enters grace window when max_turns reached.

    Spec source: spec-compliance-v2.max-turns-grace
    When max_turns is reached, enters grace window (same as budget exhaustion)
    to allow agent one final opportunity to call ctl_finish.

    Args:
        turn_count: Current turn number (already incremented by caller).
        max_turns: Maximum turns allowed (None = unlimited).
        in_grace_window: Whether already in grace window.
        budget_tracker: Budget tracker instance.
        instance_id: The agent's unique identifier.

    Returns:
        Tuple of (result_if_terminal, new_in_grace_window, max_turns_exceeded).
        result_if_terminal is None if the loop should continue.
    """
    # Check if we've exceeded max_turns
    if max_turns is None or turn_count <= max_turns:
        return None, in_grace_window, False

    # Mark that we've exceeded max_turns
    max_turns_exceeded = True

    # Enter grace window (same path as budget exhaustion)
    if not in_grace_window:
        in_grace_window = True

        # Synchronize budget tracker to grace mode
        # For max_turns case, we force entry into grace (budget may not be exhausted)
        # This sets grace_start_usage so grace budget tracking works correctly
        budget_tracker.in_grace = True
        budget_tracker.grace_start_usage = budget_tracker.used

        # Log grace entry from max_turns
        append_log(
            instance_id,
            "grace_entered",
            {
                "reason": "max_turns_exceeded",
                "turn_count": turn_count,
                "max_turns": max_turns,
                "grace_budget": budget_tracker.grace_budget,
            },
        )

        # Update state with grace mode metadata (no state change)
        update_agent_state(instance_id, {})

    # Check if grace is exhausted (handled by budget tracker)
    # Grace tokens still apply even for max_turns case
    if budget_tracker.is_grace_exhausted():
        transition_state(
            instance_id,
            LifecycleEvent.max_turns,
            error="Grace budget exhausted after max_turns without ctl_finish",
        )
        append_log(
            instance_id,
            "grace_exhausted",
            {
                "reason": "max_turns_exceeded",
                "turn_count": turn_count,
                "max_turns": max_turns,
                "grace_budget": budget_tracker.grace_budget,
            },
        )
        return (
            Failure(
                MaxTurnsReachedError(
                    f"Grace budget exhausted after max turns: {turn_count} > {max_turns}",
                    details={
                        "turn_count": turn_count,
                        "max_turns": max_turns,
                        "grace_budget": budget_tracker.grace_budget,
                    },
                )
            ),
            in_grace_window,
            max_turns_exceeded,
        )

    # Allow continuation in grace window
    return None, in_grace_window, max_turns_exceeded


__all__ = [
    "check_budget_exhaustion",
    "handle_model_request_turn",
    "track_token_usage",
    "update_state",
]
