"""Tests for token usage delta tracking (F1/F2 fix).

Verifies that PydanticAI's cumulative RunUsage is correctly delta-tracked
to avoid double-counting tokens across multiple model requests within
a single iter() call.
"""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest

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


class TestAgentRunnerDeltaTracking:
    """Test AgentRunner._track_token_usage delta computation."""

    def _make_runner(self) -> MagicMock:
        """Create a mock AgentRunner with delta tracking state."""
        from anima.shell.agent import AgentRunner

        runner = AgentRunner.__new__(AgentRunner)
        runner.token_usage = TokenUsage()
        runner.budget_tracker = BudgetTracker(budget_limit=None, grace_budget=500)
        runner.instance_id = "inst_test"
        runner.turn_count = 1
        runner._prev_run_input_tokens = 0
        runner._prev_run_output_tokens = 0
        return runner

    def test_first_usage_recorded_as_delta(self) -> None:
        """First call with cumulative usage = delta from zero."""
        runner = self._make_runner()
        usage = MagicMock()
        usage.input_tokens = 30000
        usage.output_tokens = 500

        runner._track_token_usage(usage)

        assert runner.token_usage.prompt_tokens == 30000
        assert runner.token_usage.completion_tokens == 500
        assert runner.token_usage.last_request_prompt_tokens == 30000

    def test_second_usage_only_adds_delta(self) -> None:
        """Second call should only add the delta, not the full cumulative."""
        runner = self._make_runner()

        # First call: cumulative = 30K
        usage1 = MagicMock()
        usage1.input_tokens = 30000
        usage1.output_tokens = 500
        runner._track_token_usage(usage1)

        # Second call: cumulative = 60K (delta = 30K)
        usage2 = MagicMock()
        usage2.input_tokens = 60000
        usage2.output_tokens = 1000
        runner._track_token_usage(usage2)

        # Should be 30K + 30K = 60K, NOT 30K + 60K = 90K
        assert runner.token_usage.prompt_tokens == 60000
        assert runner.token_usage.completion_tokens == 1000
        assert runner.token_usage.last_request_prompt_tokens == 60000

    def test_three_iterations_no_double_counting(self) -> None:
        """Three iter nodes should sum to exact total, not inflate."""
        runner = self._make_runner()

        for cumulative_input in [20000, 45000, 75000]:
            usage = MagicMock()
            usage.input_tokens = cumulative_input
            usage.output_tokens = cumulative_input // 10
            runner._track_token_usage(usage)

        # Total should be 75K (final cumulative), not 20K+45K+75K=140K
        assert runner.token_usage.prompt_tokens == 75000
        assert runner.token_usage.completion_tokens == 7500
        assert runner.token_usage.last_request_prompt_tokens == 75000

    def test_reset_clears_previous_snapshot(self) -> None:
        """_reset_run_usage_tracking allows fresh delta computation."""
        runner = self._make_runner()

        # First iter() call
        usage1 = MagicMock()
        usage1.input_tokens = 30000
        usage1.output_tokens = 500
        runner._track_token_usage(usage1)

        # Reset for new iter() call
        runner._reset_run_usage_tracking()

        # Second iter() call starts fresh
        usage2 = MagicMock()
        usage2.input_tokens = 40000
        usage2.output_tokens = 600
        runner._track_token_usage(usage2)

        # Total should be 30K + 40K = 70K (both iter calls)
        assert runner.token_usage.prompt_tokens == 70000
        assert runner.token_usage.completion_tokens == 1100
        # last_request reflects the latest iter's cumulative
        assert runner.token_usage.last_request_prompt_tokens == 40000

    def test_zero_delta_skipped(self) -> None:
        """No state update when usage hasn't changed."""
        runner = self._make_runner()

        usage = MagicMock()
        usage.input_tokens = 30000
        usage.output_tokens = 500
        runner._track_token_usage(usage)

        # Same usage again (no change)
        runner._track_token_usage(usage)

        assert runner.token_usage.prompt_tokens == 30000
        assert runner.token_usage.completion_tokens == 500

    def test_none_usage_ignored(self) -> None:
        """None usage is safely ignored."""
        runner = self._make_runner()
        runner._track_token_usage(None)
        assert runner.token_usage.prompt_tokens == 0


class TestInteractDeltaTracking:
    """Test process_interact_message delta tracking."""

    def test_interact_delta_not_double_counted(self) -> None:
        """Verify interact path uses delta, not cumulative.

        This is a regression test for F1: the interact path previously
        added run.usage() (cumulative) directly to token_usage on every
        node iteration, causing massive over-counting.
        """
        # This test verifies the fix exists by checking the code structure.
        # A full integration test would require a running agent.
        import inspect
        from anima.shell.agent_interact import process_interact_message

        source = inspect.getsource(process_interact_message)
        # The fix: delta computation via prev_input_tokens
        assert "prev_input_tokens" in source, (
            "process_interact_message must use delta tracking (prev_input_tokens)"
        )
        assert "delta_input" in source, (
            "process_interact_message must compute delta_input"
        )
        # The bug: direct cumulative addition should NOT be present
        assert "token_usage.prompt_tokens += prompt" not in source, (
            "process_interact_message must NOT add cumulative prompt directly"
        )
