from __future__ import annotations

import pytest
from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails
from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails

from agents import Agent, Runner
from agents.run_internal.agent_runner_helpers import snapshot_usage, usage_delta
from agents.usage import (
    RequestUsage,
    Usage,
    deserialize_usage,
    model_usage_to_span_usage,
    serialize_usage,
)
from tests.fake_model import FakeModel
from tests.test_responses import get_text_message


def test_usage_defaults_cache_write_tokens_to_zero() -> None:
    usage = Usage()

    assert usage.input_tokens_details.cached_tokens == 0
    assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 0


@pytest.mark.asyncio
async def test_runner_run_carries_request_usage_entries() -> None:
    """Ensure usage produced by the model propagates to RunResult context."""
    usage = Usage(
        requests=1,
        input_tokens=10,
        output_tokens=5,
        total_tokens=15,
        request_usage_entries=[
            RequestUsage(
                input_tokens=10,
                output_tokens=5,
                total_tokens=15,
                input_tokens_details=InputTokensDetails.model_validate(
                    {"cache_write_tokens": 0, "cached_tokens": 0}
                ),
                output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
            )
        ],
    )
    model = FakeModel(initial_output=[get_text_message("done")])
    model.set_hardcoded_usage(usage)
    agent = Agent(name="usage-agent", model=model)

    result = await Runner.run(agent, input="hi")

    propagated = result.context_wrapper.usage
    assert propagated.requests == 1
    assert propagated.total_tokens == 15
    assert len(propagated.request_usage_entries) == 1
    entry = propagated.request_usage_entries[0]
    assert entry.input_tokens == 10
    assert entry.output_tokens == 5
    assert entry.total_tokens == 15


def test_usage_add_aggregates_all_fields():
    u1 = Usage(
        requests=1,
        input_tokens=10,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 5, "cached_tokens": 3}
        ),
        output_tokens=20,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
        total_tokens=30,
    )
    u2 = Usage(
        requests=2,
        input_tokens=7,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 6, "cached_tokens": 4}
        ),
        output_tokens=8,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=6),
        total_tokens=15,
    )

    u1.add(u2)

    assert u1.requests == 3
    assert u1.input_tokens == 17
    assert u1.output_tokens == 28
    assert u1.total_tokens == 45
    assert u1.input_tokens_details.cached_tokens == 7
    assert getattr(u1.input_tokens_details, "cache_write_tokens", None) == 11
    assert u1.output_tokens_details.reasoning_tokens == 11


def test_usage_add_aggregates_with_none_values():
    u1 = Usage()
    u2 = Usage(
        requests=2,
        input_tokens=7,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 4}
        ),
        output_tokens=8,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=6),
        total_tokens=15,
    )

    u1.add(u2)

    assert u1.requests == 2
    assert u1.input_tokens == 7
    assert u1.output_tokens == 8
    assert u1.total_tokens == 15
    assert u1.input_tokens_details.cached_tokens == 4
    assert u1.output_tokens_details.reasoning_tokens == 6


def test_request_usage_creation():
    """Test that RequestUsage is created correctly."""
    request_usage = RequestUsage(
        input_tokens=100,
        output_tokens=200,
        total_tokens=300,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 10}
        ),
        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
    )

    assert request_usage.input_tokens == 100
    assert request_usage.output_tokens == 200
    assert request_usage.total_tokens == 300
    assert request_usage.input_tokens_details.cached_tokens == 10
    assert request_usage.output_tokens_details.reasoning_tokens == 20


def test_usage_add_preserves_single_request():
    """Test that adding a single request Usage creates an RequestUsage entry."""
    u1 = Usage()
    u2 = Usage(
        requests=1,
        input_tokens=100,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 10}
        ),
        output_tokens=200,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
        total_tokens=300,
    )

    u1.add(u2)

    # Should preserve the request usage details
    assert len(u1.request_usage_entries) == 1
    request_usage = u1.request_usage_entries[0]
    assert request_usage.input_tokens == 100
    assert request_usage.output_tokens == 200
    assert request_usage.total_tokens == 300
    assert request_usage.input_tokens_details.cached_tokens == 10
    assert request_usage.output_tokens_details.reasoning_tokens == 20


def test_usage_add_ignores_zero_token_requests():
    """Test that zero-token requests don't create request_usage_entries."""
    u1 = Usage()
    u2 = Usage(
        requests=1,
        input_tokens=0,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 0}
        ),
        output_tokens=0,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
        total_tokens=0,
    )

    u1.add(u2)

    # Should not create a request_usage_entry for zero tokens
    assert len(u1.request_usage_entries) == 0


def test_usage_add_ignores_multi_request_usage():
    """Test that multi-request Usage objects don't create request_usage_entries."""
    u1 = Usage()
    u2 = Usage(
        requests=3,  # Multiple requests
        input_tokens=100,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 10}
        ),
        output_tokens=200,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
        total_tokens=300,
    )

    u1.add(u2)

    # Should not create a request usage entry for multi-request usage
    assert len(u1.request_usage_entries) == 0


def test_usage_add_merges_existing_request_usage_entries():
    """Test that existing request_usage_entries are merged when adding Usage objects."""
    # Create first usage with request_usage_entries
    u1 = Usage()
    u2 = Usage(
        requests=1,
        input_tokens=100,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 10}
        ),
        output_tokens=200,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
        total_tokens=300,
    )
    u1.add(u2)

    # Create second usage with request_usage_entries
    u3 = Usage(
        requests=1,
        input_tokens=50,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 5}
        ),
        output_tokens=75,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
        total_tokens=125,
    )

    u1.add(u3)

    # Should have both request_usage_entries
    assert len(u1.request_usage_entries) == 2

    # First request
    first = u1.request_usage_entries[0]
    assert first.input_tokens == 100
    assert first.output_tokens == 200
    assert first.total_tokens == 300

    # Second request
    second = u1.request_usage_entries[1]
    assert second.input_tokens == 50
    assert second.output_tokens == 75
    assert second.total_tokens == 125


def test_usage_add_with_pre_existing_request_usage_entries():
    """Test adding Usage objects that already have request_usage_entries."""
    u1 = Usage()

    # Create a usage with request_usage_entries
    u2 = Usage(
        requests=1,
        input_tokens=100,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 10}
        ),
        output_tokens=200,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=20),
        total_tokens=300,
    )
    u1.add(u2)

    # Create another usage with request_usage_entries
    u3 = Usage(
        requests=1,
        input_tokens=50,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 5}
        ),
        output_tokens=75,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
        total_tokens=125,
    )

    # Add u3 to u1
    u1.add(u3)

    # Should have both request_usage_entries
    assert len(u1.request_usage_entries) == 2
    assert u1.request_usage_entries[0].input_tokens == 100
    assert u1.request_usage_entries[1].input_tokens == 50


def test_usage_add_preserves_existing_entries_when_top_level_also_set():
    """When `other` has both top-level single-request fields AND pre-populated
    `request_usage_entries`, the existing entries (which carry the authoritative
    nested token details) must not be discarded in favor of a synthesized entry
    built from only the top-level fields.
    """
    u1 = Usage()
    u2 = Usage(
        requests=1,
        input_tokens=100,
        output_tokens=50,
        total_tokens=150,
        request_usage_entries=[
            RequestUsage(
                input_tokens=100,
                output_tokens=50,
                total_tokens=150,
                input_tokens_details=InputTokensDetails.model_validate(
                    {"cache_write_tokens": 0, "cached_tokens": 10}
                ),
                output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
            )
        ],
    )

    u1.add(u2)

    # The pre-populated entry must be preserved — including its nested details —
    # rather than being replaced by a synthesized entry with zeroed-out details.
    assert len(u1.request_usage_entries) == 1
    entry = u1.request_usage_entries[0]
    assert entry.input_tokens_details.cached_tokens == 10
    assert entry.output_tokens_details.reasoning_tokens == 5


def test_usage_request_usage_entries_default_empty():
    """Test that request_usage_entries defaults to an empty list."""
    u = Usage()
    assert u.request_usage_entries == []


def test_anthropic_cost_calculation_scenario():
    """Test a realistic scenario for Sonnet 4.5 cost calculation with 200K token thresholds."""
    # Simulate 3 API calls: 100K, 150K, and 80K input tokens each
    # None exceed 200K, so they should all use the lower pricing tier

    usage = Usage()

    # First request: 100K input tokens
    req1 = Usage(
        requests=1,
        input_tokens=100_000,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 0}
        ),
        output_tokens=50_000,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
        total_tokens=150_000,
    )
    usage.add(req1)

    # Second request: 150K input tokens
    req2 = Usage(
        requests=1,
        input_tokens=150_000,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 0}
        ),
        output_tokens=75_000,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
        total_tokens=225_000,
    )
    usage.add(req2)

    # Third request: 80K input tokens
    req3 = Usage(
        requests=1,
        input_tokens=80_000,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 0, "cached_tokens": 0}
        ),
        output_tokens=40_000,
        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
        total_tokens=120_000,
    )
    usage.add(req3)

    # Verify aggregated totals
    assert usage.requests == 3
    assert usage.input_tokens == 330_000  # 100K + 150K + 80K
    assert usage.output_tokens == 165_000  # 50K + 75K + 40K
    assert usage.total_tokens == 495_000  # 150K + 225K + 120K

    # Verify request_usage_entries preservation
    assert len(usage.request_usage_entries) == 3
    assert usage.request_usage_entries[0].input_tokens == 100_000
    assert usage.request_usage_entries[1].input_tokens == 150_000
    assert usage.request_usage_entries[2].input_tokens == 80_000

    # All request_usage_entries are under 200K threshold
    for req in usage.request_usage_entries:
        assert req.input_tokens < 200_000
        assert req.output_tokens < 200_000


def test_usage_normalizes_none_token_details():
    # Some providers don't populate optional token detail fields
    # (cached_tokens, reasoning_tokens), and the OpenAI SDK's generated
    # code can bypass Pydantic validation (e.g., via model_construct),
    # allowing None values. We normalize these to 0 to prevent TypeErrors.

    # Test entire objects being None (BeforeValidator)
    usage = Usage(
        requests=1,
        input_tokens=100,
        input_tokens_details=None,  # type: ignore[arg-type]
        output_tokens=50,
        output_tokens_details=None,  # type: ignore[arg-type]
        total_tokens=150,
    )
    assert usage.input_tokens_details.cached_tokens == 0
    assert usage.output_tokens_details.reasoning_tokens == 0

    # Test fields within objects being None (__post_init__)
    input_details = InputTokensDetails.model_validate({"cache_write_tokens": 0, "cached_tokens": 0})
    input_details.__dict__["cached_tokens"] = None
    input_details.__dict__["cache_write_tokens"] = None

    output_details = OutputTokensDetails(reasoning_tokens=0)
    output_details.__dict__["reasoning_tokens"] = None

    usage = Usage(
        requests=1,
        input_tokens=100,
        input_tokens_details=input_details,
        output_tokens=50,
        output_tokens_details=output_details,
        total_tokens=150,
    )

    # __post_init__ should normalize None to 0
    assert usage.input_tokens_details.cached_tokens == 0
    assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 0
    assert usage.output_tokens_details.reasoning_tokens == 0


def test_usage_normalizes_chat_completions_types():
    # Chat Completions API uses PromptTokensDetails and CompletionTokensDetails,
    # while Usage expects InputTokensDetails and OutputTokensDetails (Responses API).
    # The BeforeValidator should convert between these types.

    prompt_details = PromptTokensDetails.model_validate(
        {
            "audio_tokens": 10,
            "cached_tokens": 50,
            "cache_write_tokens": 7,
        }
    )
    completion_details = CompletionTokensDetails(
        accepted_prediction_tokens=5,
        audio_tokens=10,
        reasoning_tokens=100,
        rejected_prediction_tokens=2,
    )

    usage = Usage(
        requests=1,
        input_tokens=200,
        input_tokens_details=prompt_details,  # type: ignore[arg-type]
        output_tokens=150,
        output_tokens_details=completion_details,  # type: ignore[arg-type]
        total_tokens=350,
    )

    # Should convert to Responses API types, extracting the relevant fields
    assert isinstance(usage.input_tokens_details, InputTokensDetails)
    assert usage.input_tokens_details.cached_tokens == 50
    assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 7

    assert isinstance(usage.output_tokens_details, OutputTokensDetails)
    assert usage.output_tokens_details.reasoning_tokens == 100


def test_usage_serialization_preserves_cache_write_tokens() -> None:
    usage = Usage(
        requests=1,
        input_tokens=20,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 7, "cached_tokens": 3}
        ),
        output_tokens=5,
        total_tokens=25,
        request_usage_entries=[
            RequestUsage(
                input_tokens=20,
                output_tokens=5,
                total_tokens=25,
                input_tokens_details=InputTokensDetails.model_validate(
                    {"cache_write_tokens": 7, "cached_tokens": 3}
                ),
                output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
            )
        ],
    )

    serialized = serialize_usage(usage)
    restored = deserialize_usage(serialized)

    assert serialized["input_tokens_details"] == [{"cached_tokens": 3, "cache_write_tokens": 7}]
    assert getattr(restored.input_tokens_details, "cache_write_tokens", None) == 7
    assert (
        getattr(
            restored.request_usage_entries[0].input_tokens_details,
            "cache_write_tokens",
            None,
        )
        == 7
    )


def test_usage_deserialization_defaults_legacy_cache_write_tokens() -> None:
    restored = deserialize_usage(
        {
            "requests": 1,
            "input_tokens": 20,
            "output_tokens": 5,
            "total_tokens": 25,
            "input_tokens_details": [{"cached_tokens": 3}],
            "request_usage_entries": [
                {
                    "input_tokens": 20,
                    "output_tokens": 5,
                    "total_tokens": 25,
                    "input_tokens_details": {"cached_tokens": 3},
                }
            ],
        }
    )

    assert restored.input_tokens_details.cached_tokens == 3
    assert getattr(restored.input_tokens_details, "cache_write_tokens", None) == 0
    assert restored.request_usage_entries[0].input_tokens_details.cached_tokens == 3
    assert (
        getattr(
            restored.request_usage_entries[0].input_tokens_details,
            "cache_write_tokens",
            None,
        )
        == 0
    )


def test_usage_snapshot_delta_and_span_preserve_cache_write_tokens() -> None:
    start = Usage(
        requests=1,
        input_tokens=10,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 2, "cached_tokens": 3}
        ),
        output_tokens=4,
        total_tokens=14,
    )
    end = Usage(
        requests=2,
        input_tokens=30,
        input_tokens_details=InputTokensDetails.model_validate(
            {"cache_write_tokens": 9, "cached_tokens": 8}
        ),
        output_tokens=10,
        total_tokens=40,
    )

    snapshot = snapshot_usage(start)
    delta = usage_delta(snapshot, end)
    span_usage = model_usage_to_span_usage(delta)

    assert getattr(snapshot.input_tokens_details, "cache_write_tokens", None) == 2
    assert delta.input_tokens_details.cached_tokens == 5
    assert getattr(delta.input_tokens_details, "cache_write_tokens", None) == 7
    assert span_usage["input_tokens_details"] == {
        "cached_tokens": 5,
        "cache_write_tokens": 7,
    }