"""Tests for amplitude_ai.utils.costs."""
from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

import pytest

_MOD = "amplitude_ai.utils.costs"


class TestGetGenaiPriceLookupCandidates:
    """Tests for provider/model candidate generation."""

    def test_bare_model_no_provider(self) -> None:
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates("gpt-4o")
        assert candidates == [("gpt-4o", None)]

    def test_prefixed_model_extracts_provider(self) -> None:
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates("openai:gpt-4o")
        assert candidates == [("gpt-4o", "openai")]

    def test_bedrock_prefix_maps_to_aws(self) -> None:
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates("bedrock:anthropic.claude-sonnet-4-6")
        model_refs = [c[0] for c in candidates]
        providers = [c[1] for c in candidates]
        assert "anthropic.claude-sonnet-4-6" in model_refs
        assert "regional.anthropic.claude-sonnet-4-6" in model_refs
        assert "global.anthropic.claude-sonnet-4-6" in model_refs
        assert ("claude-sonnet-4-6", None) in candidates
        assert all(p in ("aws", None) for p in providers)

    def test_default_provider_used_for_bare_bedrock_model(self) -> None:
        """When model lacks prefix, default_provider supplies the provider."""
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates(
            "anthropic.claude-sonnet-4-6", default_provider="bedrock"
        )
        model_refs = [c[0] for c in candidates]
        assert "anthropic.claude-sonnet-4-6" in model_refs
        assert "regional.anthropic.claude-sonnet-4-6" in model_refs
        assert ("claude-sonnet-4-6", None) in candidates

    def test_bedrock_region_prefix_stripped(self) -> None:
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates("bedrock:us.anthropic.claude-sonnet-4-6")
        model_refs = [c[0] for c in candidates]
        assert "us.anthropic.claude-sonnet-4-6" in model_refs
        assert "anthropic.claude-sonnet-4-6" in model_refs
        assert "claude-sonnet-4-6" in model_refs

    def test_bedrock_apac_region_prefix_stripped(self) -> None:
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates(
            "apac.anthropic.claude-sonnet-4-6", default_provider="bedrock"
        )
        model_refs = [c[0] for c in candidates]
        assert "anthropic.claude-sonnet-4-6" in model_refs
        assert "global.anthropic.claude-sonnet-4-6" in model_refs

    def test_unknown_region_and_vendor_handled_without_code_changes(self) -> None:
        """Generalized dot-stripping handles future regions/vendors automatically."""
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates(
            "sa-east-1.newvendor.some-model-v3", default_provider="bedrock"
        )
        model_refs = [c[0] for c in candidates]
        assert "newvendor.some-model-v3" in model_refs
        assert "some-model-v3" in model_refs

    def test_two_segment_bedrock_model_gets_regional_global(self) -> None:
        """Two-segment vendor.model gets regional./global. prefixes."""
        from amplitude_ai.utils.costs import get_genai_price_lookup_candidates

        candidates = get_genai_price_lookup_candidates(
            "anthropic.claude-sonnet-4-6", default_provider="bedrock"
        )
        model_refs = [c[0] for c in candidates]
        assert "regional.anthropic.claude-sonnet-4-6" in model_refs
        assert "global.anthropic.claude-sonnet-4-6" in model_refs


class TestCalculateCost:
    """Tests for calculate_cost with default_provider."""

    @patch(f"{_MOD}.GENAI_PRICES_AVAILABLE", True)
    @patch(f"{_MOD}.genai_calc_price")
    @patch(f"{_MOD}.GenAIUsage")
    def test_default_provider_forwarded_to_candidates(
        self, mock_usage_cls: MagicMock, mock_calc: MagicMock
    ) -> None:
        """default_provider should be threaded through to genai_calc_price."""
        from amplitude_ai.utils.costs import calculate_cost

        mock_price = MagicMock()
        mock_price.total_price = 0.005
        mock_calc.return_value = mock_price

        result = calculate_cost(
            model_name="anthropic.claude-sonnet-4-6",
            input_tokens=1000,
            output_tokens=100,
            default_provider="bedrock",
        )

        assert result == 0.005
        first_call_kwargs = mock_calc.call_args_list[0]
        assert first_call_kwargs.kwargs.get("provider_id") == "aws" or \
            first_call_kwargs[1].get("provider_id") == "aws"

    @patch(f"{_MOD}.GENAI_PRICES_AVAILABLE", True)
    @patch(f"{_MOD}.genai_calc_price")
    @patch(f"{_MOD}.GenAIUsage")
    def test_no_default_provider(
        self, mock_usage_cls: MagicMock, mock_calc: MagicMock
    ) -> None:
        """Without default_provider and no prefix, provider_id should be None."""
        from amplitude_ai.utils.costs import calculate_cost

        mock_price = MagicMock()
        mock_price.total_price = 0.003
        mock_calc.return_value = mock_price

        result = calculate_cost(
            model_name="anthropic.claude-sonnet-4-6",
            input_tokens=1000,
            output_tokens=100,
        )

        assert result == 0.003
        first_call_kwargs = mock_calc.call_args_list[0]
        assert first_call_kwargs.kwargs.get("provider_id") is None or \
            first_call_kwargs[1].get("provider_id") is None

    @patch(f"{_MOD}.GENAI_PRICES_AVAILABLE", True)
    @patch(f"{_MOD}.genai_calc_price")
    @patch(f"{_MOD}.GenAIUsage")
    def test_explicit_prefix_overrides_default_provider(
        self, mock_usage_cls: MagicMock, mock_calc: MagicMock
    ) -> None:
        """A provider: prefix in model_name takes precedence over default_provider."""
        from amplitude_ai.utils.costs import calculate_cost

        mock_price = MagicMock()
        mock_price.total_price = 0.004
        mock_calc.return_value = mock_price

        result = calculate_cost(
            model_name="openai:gpt-4o",
            input_tokens=1000,
            output_tokens=100,
            default_provider="bedrock",
        )

        assert result == 0.004
        first_call_kwargs = mock_calc.call_args_list[0]
        assert first_call_kwargs.kwargs.get("provider_id") == "openai" or \
            first_call_kwargs[1].get("provider_id") == "openai"

    @pytest.mark.integration()
    def test_bedrock_model_with_default_provider_returns_nonzero(self) -> None:
        """Integration: genai-prices should resolve Bedrock model with default_provider."""
        from amplitude_ai.utils.costs import GENAI_PRICES_AVAILABLE, calculate_cost

        if not GENAI_PRICES_AVAILABLE:
            pytest.skip("genai-prices not installed")

        cost = calculate_cost(
            model_name="anthropic.claude-sonnet-4-20250514-v1:0",
            input_tokens=100_000,
            output_tokens=10_000,
            default_provider="bedrock",
        )
        assert cost > 0, "Bedrock model with default_provider='bedrock' should produce nonzero cost"

    @pytest.mark.integration()
    def test_bedrock_model_without_provider_still_resolves_via_dot_stripping(self) -> None:
        """Integration: even without default_provider, generalized dot-stripping
        finds the bare model name which genai-prices resolves via auto-detect."""
        from amplitude_ai.utils.costs import GENAI_PRICES_AVAILABLE, calculate_cost

        if not GENAI_PRICES_AVAILABLE:
            pytest.skip("genai-prices not installed")

        cost = calculate_cost(
            model_name="anthropic.claude-sonnet-4-20250514-v1:0",
            input_tokens=100_000,
            output_tokens=10_000,
        )
        assert cost > 0, "Dot-stripping should resolve vendor.model even without default_provider"

    @pytest.mark.integration()
    def test_bedrock_region_prefixed_model_resolves(self) -> None:
        """Integration: us.anthropic.X with bedrock provider should resolve."""
        from amplitude_ai.utils.costs import GENAI_PRICES_AVAILABLE, calculate_cost

        if not GENAI_PRICES_AVAILABLE:
            pytest.skip("genai-prices not installed")

        cost = calculate_cost(
            model_name="us.anthropic.claude-sonnet-4-20250514-v1:0",
            input_tokens=100_000,
            output_tokens=10_000,
            default_provider="bedrock",
        )
        assert cost > 0, "Region-prefixed Bedrock model should produce nonzero cost"


class TestAmplitudeAIAutoCalculateCost:
    """Tests for AmplitudeAI.track_ai_message auto-cost with default_provider."""

    @patch("amplitude_ai.client.track_ai_message")
    @patch("amplitude_ai.client.calculate_cost", return_value=0.42)
    def test_auto_cost_passes_provider_as_default(
        self, mock_cost: MagicMock, mock_track: MagicMock
    ) -> None:
        """When total_cost_usd is None, calculate_cost should receive provider as default_provider."""
        from amplitude_ai.client import AmplitudeAI

        mock_track.return_value = "msg-123"
        ai = AmplitudeAI(amplitude=MagicMock())

        ai.track_ai_message(
            user_id="u1",
            content="Hello",
            session_id="s1",
            model="anthropic.claude-sonnet-4-6",
            provider="bedrock",
            latency_ms=100,
            input_tokens=1000,
            output_tokens=100,
            total_cost_usd=None,
        )

        mock_cost.assert_called_once()
        kwargs = mock_cost.call_args.kwargs
        assert kwargs["default_provider"] == "bedrock"
        assert kwargs["model_name"] == "anthropic.claude-sonnet-4-6"

    @patch("amplitude_ai.client.track_ai_message")
    @patch("amplitude_ai.client.calculate_cost", return_value=0.42)
    def test_auto_cost_skipped_when_cost_provided(
        self, mock_cost: MagicMock, mock_track: MagicMock
    ) -> None:
        """When total_cost_usd is explicitly provided, calculate_cost should not be called."""
        from amplitude_ai.client import AmplitudeAI

        mock_track.return_value = "msg-123"
        ai = AmplitudeAI(amplitude=MagicMock())

        ai.track_ai_message(
            user_id="u1",
            content="Hello",
            session_id="s1",
            model="gpt-4o",
            provider="openai",
            latency_ms=100,
            input_tokens=1000,
            output_tokens=100,
            total_cost_usd=0.05,
        )

        mock_cost.assert_not_called()
