"""Targeted tests for budget exhaustion error type compliance.

Spec source: spec-compliance-v2.budget-error-type requires all budget
exhaustion paths return Failure(BudgetExhaustedError code 310), not
success-failed result or internal error.

Tests cover:
1. Main loop budget exhaustion (via _process_model_result)
2. Interact loop budget exhaustion (via _process_interact_message)
3. Grace exhaustion without ctl_finish
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, cast

import pytest
from returns.result import Failure, Success

from anima.core.lifecycle import LifecycleEvent
from anima.core.schemas import (
    AgentState,
    BudgetExhaustedError,
    CompletionPolicy,
    JobSpec,
    LifecycleType,
)
from anima.shell.agent import AgentRunner
from anima.shell.state import (
    create_agent_state,
    read_state,
    transition_state,
)


def _create_running_runner(
    *,
    instance_id: str,
    max_turns: int | None = None,
    lifecycle: LifecycleType = LifecycleType.ephemeral,
    max_tokens: int | None = None,
    grace_budget: int = 0,
) -> AgentRunner:
    """Create an AgentRunner with persisted running state."""
    job = JobSpec(
        persona_ref="test@test",
        instance_id=instance_id,
        max_turns=max_turns,
        max_tokens=max_tokens,
        finalize_grace_budget=grace_budget,
        lifecycle=lifecycle,
    )
    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)
    return AgentRunner(
        instance_id=instance_id,
        task="test task",
        job=job,
        completion_policy=CompletionPolicy.auto,
        max_turns=max_turns,
    )


# =============================================================================
# Test 1: Main Loop Budget Exhaustion
# =============================================================================


class TestMainLoopBudgetExhaustion:
    """Budget exhaustion in main agent loop via _process_model_result."""

    def test_process_model_result_returns_budget_error_on_exhaustion(
        self, tmp_state_dir: Any
    ) -> None:
        """_process_model_result returns Failure(BudgetExhaustedError) when budget exhausted."""
        _ = tmp_state_dir
        from pydantic_ai._agent_graph import CallToolsNode, ModelRequestNode
        from pydantic_ai.messages import ModelResponse, TextPart

        runner = _create_running_runner(
            instance_id="inst_main_budget",
            max_turns=10,
            max_tokens=10,
            grace_budget=0,
        )

        # Simulate budget exhaustion
        runner.budget_tracker.record_usage(10)

        response = ModelResponse(parts=[TextPart(content="done")])
        node = CallToolsNode(model_response=response, tool_call_results={})

        run = _FakeIterRun(
            nodes=[node],
            messages=[response],
            usage=_UsageSnapshot(5, 5),  # Additional usage pushes over budget
        )

        result, has_usage = runner._process_model_result(
            node, run, ModelRequestNode, CallToolsNode
        )

        assert result is not None
        assert isinstance(result, Failure)
        error = result.failure()
        assert isinstance(error, BudgetExhaustedError)
        assert error.code == 310
        assert "Budget exhausted" in error.message


# =============================================================================
# Test 2: Interact Loop Budget Exhaustion (Missing Check Bug)
# =============================================================================


class TestInteractLoopBudgetExhaustion:
    """Budget exhaustion in interact message loop via _process_interact_message."""

    @pytest.mark.asyncio
    async def test_interact_loop_returns_budget_error_on_exhaustion(
        self, tmp_state_dir: Any
    ) -> None:
        """_process_interact_message returns Failure(BudgetExhaustedError) when budget exhausted during interact.

        This is the missing check bug: interact loop tracks usage but doesn't check exhaustion.
        After fix, this should return Failure(BudgetExhaustedError).
        """
        _ = tmp_state_dir
        from pydantic_ai._agent_graph import CallToolsNode, ModelRequestNode
        from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart

        runner = _create_running_runner(
            instance_id="inst_interact_budget",
            max_turns=10,
            max_tokens=10,
            grace_budget=0,
        )

        from tests.helpers.mcp_fakes import FakeMCPServer

        fake_mcp = FakeMCPServer()
        runner.mcp_server = fake_mcp

        # Exhaust budget before interact
        runner.budget_tracker.record_usage(10)

        # Create nodes that add more token usage
        request = ModelRequest(parts=[])
        response = ModelResponse(parts=[TextPart(content="response")])
        nodes = [
            ModelRequestNode(request=request),
            CallToolsNode(model_response=response, tool_call_results={}),
        ]
        history = [response]
        run = _FakeIterRun(nodes=nodes, messages=history, usage=_UsageSnapshot(5, 5))
        agent = _FakeAgent(run)

        result = await runner._process_interact_message(
            cast(Any, agent), "test message", []
        )

        # After fix: Should return Failure(BudgetExhaustedError)
        assert isinstance(result, Failure)
        error = result.failure()
        assert isinstance(error, BudgetExhaustedError)
        assert error.code == 310

        state_result = read_state(runner.instance_id)
        assert isinstance(state_result, Success)
        assert state_result.unwrap().state == AgentState.failed


# =============================================================================
# Test 3: Grace Exhaustion Without Finish
# =============================================================================


class TestGraceExhaustion:
    """Grace budget exhausted without ctl_finish."""

    def test_grace_exhaustion_returns_budget_error(self, tmp_state_dir: Any) -> None:
        """Grace exhaustion returns Failure(BudgetExhaustedError code 310)."""
        _ = tmp_state_dir
        job = JobSpec(
            persona_ref="test@test",
            instance_id="inst_grace_exhaust",
            max_tokens=10,
            finalize_grace_budget=5,  # Grace available
        )
        created = create_agent_state("inst_grace_exhaust", "test-persona", job)
        assert isinstance(created, Success)
        instance_id = created.unwrap().instance_id
        started = transition_state(instance_id, LifecycleEvent.start)
        assert isinstance(started, Success)

        runner = AgentRunner(instance_id=instance_id, task="task", job=job)

        # Exhaust main budget, enter grace
        runner.budget_tracker.record_usage(10)
        first = runner._check_budget_exhaustion()
        assert first is None  # Enters grace window
        assert runner.in_grace_window is True

        # Exhaust grace budget without ctl_finish
        runner.budget_tracker.record_usage(5)
        result = runner._check_budget_exhaustion()

        # Must return Failure(BudgetExhaustedError)
        assert isinstance(result, Failure)
        error = result.failure()
        assert isinstance(error, BudgetExhaustedError)
        assert error.code == 310
        assert "grace" in error.message.lower()

        state_result = read_state(instance_id)
        assert isinstance(state_result, Success)
        assert state_result.unwrap().state == AgentState.failed


# =============================================================================
# Test 4: All Exhaustion Paths Return Same Error Type
# =============================================================================


class TestExhaustionErrorTypeConsistency:
    """All budget exhaustion paths return BudgetExhaustedError code 310."""

    def test_all_paths_return_budgetexhaustederror_code_310(
        self, tmp_state_dir: Any
    ) -> None:
        """Verify all exhaustion paths use BudgetExhaustedError code 310."""
        _ = tmp_state_dir

        # Path 1: No grace budget, direct exhaustion
        job1 = JobSpec(
            persona_ref="test@test",
            instance_id="inst_no_grace",
            max_tokens=5,
            finalize_grace_budget=0,
        )
        create_agent_state("inst_no_grace", "test-persona", job1)
        transition_state("inst_no_grace", LifecycleEvent.start)
        runner1 = AgentRunner(instance_id="inst_no_grace", task="task", job=job1)
        runner1.budget_tracker.record_usage(5)
        result1 = runner1._check_budget_exhaustion()

        assert isinstance(result1, Failure)
        error1 = result1.failure()
        assert isinstance(error1, BudgetExhaustedError)
        assert error1.code == 310  # SPEC-COMPLIANCE-V2

        # Path 2: Grace exhaustion
        job2 = JobSpec(
            persona_ref="test@test",
            instance_id="inst_grace",
            max_tokens=5,
            finalize_grace_budget=3,
        )
        create_agent_state("inst_grace", "test-persona", job2)
        transition_state("inst_grace", LifecycleEvent.start)
        runner2 = AgentRunner(instance_id="inst_grace", task="task", job=job2)
        runner2.budget_tracker.record_usage(5)
        runner2._check_budget_exhaustion()  # Enter grace
        runner2.budget_tracker.record_usage(3)
        result2 = runner2._check_budget_exhaustion()

        assert isinstance(result2, Failure)
        error2 = result2.failure()
        assert isinstance(error2, BudgetExhaustedError)
        assert error2.code == 310  # SPEC-COMPLIANCE-V2


# =============================================================================
# Helpers
# =============================================================================


@dataclass
class _Usage:
    input_tokens: int
    output_tokens: int


class _UsageSnapshot:
    def __init__(self, input_tokens: int, output_tokens: int) -> None:
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens


class _FakeIterRun:
    def __init__(
        self, nodes: list[Any], messages: list[Any], usage: _UsageSnapshot
    ) -> None:
        self._nodes = nodes
        self._messages = messages
        self._usage = usage
        self._index = 0

    def __aiter__(self) -> "_FakeIterRun":
        self._index = 0
        return self

    async def __anext__(self) -> Any:
        if self._index >= len(self._nodes):
            raise StopAsyncIteration
        node = self._nodes[self._index]
        self._index += 1
        return node

    def usage(self) -> _UsageSnapshot:
        return self._usage

    def all_messages(self) -> list[Any]:
        return self._messages


class _FakeIterContext:
    def __init__(self, run: _FakeIterRun) -> None:
        self._run = run

    async def __aenter__(self) -> _FakeIterRun:
        return self._run

    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
        return False


class _FakeAgent:
    def __init__(self, run: _FakeIterRun) -> None:
        self._run = run

    def iter(self, **_: Any) -> _FakeIterContext:
        return _FakeIterContext(self._run)
