"""Test max_turns grace window behavior.

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

from __future__ import annotations

from typing import Any

from returns.result import Failure, Success

from anima.core.lifecycle import LifecycleEvent
from anima.core.schemas import JobSpec
from anima.shell.agent import AgentRunner
from anima.shell.state import (
    create_agent_state,
    transition_state,
)


def _create_running_runner(
    instance_id: str,
    max_turns: int | None = None,
    finalize_grace_budget: int = 500,
) -> AgentRunner:
    """Create a running AgentRunner for testing."""
    job = JobSpec(
        persona_ref="test@test",
        instance_id=instance_id,
        max_turns=max_turns,
        finalize_grace_budget=finalize_grace_budget,
    )
    created = create_agent_state(instance_id, "test-persona", job)
    assert isinstance(created, Success)
    instance_id = created.unwrap().instance_id
    started = transition_state(instance_id, LifecycleEvent.start)
    assert isinstance(started, Success)

    # IMPORTANT: AgentRunner constructor requires max_turns as explicit parameter
    # It does NOT pull from job.max_turns
    runner = AgentRunner(
        instance_id=instance_id, task="task", job=job, max_turns=max_turns
    )
    return runner


def test_max_turns_enters_grace_window(tmp_state_dir: Any) -> None:
    """When max_turns reached, agent enters grace window (not immediate failure)."""
    _ = tmp_state_dir
    runner = _create_running_runner(
        instance_id="inst_max_turns_grace",
        max_turns=2,
        finalize_grace_budget=100,
    )

    # At turn 1, within limit (turn_count becomes 2 after increment)
    runner.turn_count = 1
    result = runner._handle_model_request()
    assert result is None
    assert runner.in_grace_window is False
    # max_turns_exceeded is set when turn_count > max_turns, not before
    assert runner.turn_count == 2  # Incremented to 2

    # At turn 2, exceeds max_turns -> enters grace (turn_count becomes 3 after increment)
    runner.turn_count = 2
    result = runner._handle_model_request()
    assert result is None  # NOT immediate failure
    assert runner.in_grace_window is True
    assert runner.max_turns_exceeded is True
    assert runner.turn_count == 3  # Incremented to 3


def test_max_turns_grace_allows_finish(tmp_state_dir: Any) -> None:
    """Agent in max_turns grace can call ctl_finish."""
    _ = tmp_state_dir
    runner = _create_running_runner(
        instance_id="inst_max_turns_finish",
        max_turns=1,
        finalize_grace_budget=100,
    )

    # Exceed max_turns -> grace
    runner.turn_count = 2
    result = runner._handle_model_request()
    assert result is None
    assert runner.in_grace_window is True

    # Grace window active - can continue (would call ctl_finish)
    runner.turn_count = 3
    result = runner._handle_model_request()
    # Still in grace, can finish
    assert result is None


def test_max_turns_grace_exhaustion_fails(tmp_state_dir: Any) -> None:
    """Grace budget exhaustion after max_turns returns Failure."""
    _ = tmp_state_dir
    runner = _create_running_runner(
        instance_id="inst_max_turns_exhaust",
        max_turns=1,
        finalize_grace_budget=10,  # Small grace budget
    )

    # Exceed max_turns -> grace
    runner.turn_count = 2
    result = runner._handle_model_request()
    assert result is None
    assert runner.in_grace_window is True

    # Use all grace budget
    runner.budget_tracker.record_usage(100)  # Exhaust budget
    runner.budget_tracker.enter_grace()
    runner.budget_tracker.record_usage(10)  # Exhaust grace

    # Now grace is exhausted -> Failure
    runner.turn_count = 3
    result = runner._handle_model_request()
    assert isinstance(result, Failure)
    assert "Grace budget exhausted" in str(result.failure())
    assert runner.grace_exhausted_without_finish is True


def test_max_turns_zero_grace_budget_fails_immediately(tmp_state_dir: Any) -> None:
    """Zero grace_budget produces immediate failure on max_turns (no grace)."""
    _ = tmp_state_dir
    runner = _create_running_runner(
        instance_id="inst_max_turns_no_grace",
        max_turns=1,
        finalize_grace_budget=0,  # No grace
    )

    # Exceed max_turns
    runner.turn_count = 2
    result = runner._handle_model_request()

    # Grace exhausted immediately (grace_budget=0)
    assert isinstance(result, Failure)
    assert "Grace budget exhausted" in str(result.failure())


def test_max_turns_grace_message_injection(tmp_state_dir: Any) -> None:
    """Max_turns grace enters same grace window as budget exhaustion."""
    _ = tmp_state_dir
    runner = _create_running_runner(
        instance_id="inst_max_turns_grace_msg",
        max_turns=2,
        finalize_grace_budget=100,
    )

    # Verify budget tracker has grace budget
    assert runner.budget_tracker.grace_budget == 100

    # Exceed max_turns
    runner.turn_count = 3
    result = runner._handle_model_request()
    assert result is None
    assert runner.in_grace_window is True

    # Grace window state is set (same as budget exhaustion path)
    assert runner.grace_message_injected is False  # Will be injected on next turn
