"""Tests for cost/pricing calculations."""

import pytest

from qwen_pipeline_server import OMNI_PRICING, OMNI_MODEL


class TestOmniPricing:
    """Tests for OMNI_PRICING rates and cost calculations."""

    def test_plus_pricing_exists(self):
        assert 'qwen3.5-omni-plus-realtime' in OMNI_PRICING

    def test_flash_pricing_exists(self):
        assert 'qwen3.5-omni-flash-realtime' in OMNI_PRICING

    def test_plus_rates(self):
        p = OMNI_PRICING['qwen3.5-omni-plus-realtime']
        assert p['input_text'] == 2.1
        assert p['input_audio'] == 16.5
        assert p['output_audio'] == 62.0

    def test_flash_rates(self):
        p = OMNI_PRICING['qwen3.5-omni-flash-realtime']
        assert p['input_text'] == 0.55
        assert p['input_audio'] == 4.5
        assert p['output_audio'] == 17.7

    def test_no_output_text_in_pricing(self):
        """Output text is free when audio output is present."""
        for model, p in OMNI_PRICING.items():
            assert 'output_text' not in p

    def test_plus_cost_calculation(self):
        """Real example: 277 input text + 505 input audio + 124 output audio."""
        p = OMNI_PRICING['qwen3.5-omni-plus-realtime']
        cost = (
            277 * p['input_text'] / 1_000_000
            + 505 * p['input_audio'] / 1_000_000
            + 124 * p['output_audio'] / 1_000_000
        )
        # (277*2.1 + 505*16.5 + 124*62.0) / 1_000_000 = 16602.2 / 1_000_000
        expected = (277 * 2.1 + 505 * 16.5 + 124 * 62.0) / 1_000_000
        assert abs(cost - expected) < 1e-10
        assert cost > 0

    def test_flash_cost_calculation(self):
        """Same tokens with flash model should be cheaper."""
        plus = OMNI_PRICING['qwen3.5-omni-plus-realtime']
        flash = OMNI_PRICING['qwen3.5-omni-flash-realtime']
        tokens = {'input_text': 277, 'input_audio': 505, 'output_audio': 124}

        plus_cost = sum(
            tokens[k] * plus[k] / 1_000_000 for k in tokens)
        flash_cost = sum(
            tokens[k] * flash[k] / 1_000_000 for k in tokens)
        assert flash_cost < plus_cost

    def test_zero_tokens(self):
        p = OMNI_PRICING['qwen3.5-omni-plus-realtime']
        cost = (
            0 * p['input_text'] / 1_000_000
            + 0 * p['input_audio'] / 1_000_000
            + 0 * p['output_audio'] / 1_000_000
        )
        assert cost == 0.0

    def test_unknown_model_fallback(self):
        """Unknown model should fall back to OMNI_MODEL (plus) pricing."""
        p = OMNI_PRICING.get('nonexistent-model', OMNI_PRICING[OMNI_MODEL])
        assert p == OMNI_PRICING['qwen3.5-omni-plus-realtime']

    def test_default_omni_model_is_plus(self):
        assert OMNI_MODEL == 'qwen3.5-omni-plus-realtime'