from contextlib import contextmanager
from datetime import datetime, timezone
from decimal import Decimal
from pathlib import Path
from types import SimpleNamespace
import unittest
from unittest.mock import patch
from uuid import uuid4

from core.capabilities import ModelCapabilities
from core.pricing import calculate_chat_cost, resolve_chat_price
from core.storage.usage import record_chat_usage


MODELS_DIR = Path(__file__).resolve().parents[1] / "config" / "models"


class PricingTests(unittest.TestCase):
    def test_deepseek_flash_selects_peak_and_off_peak(self) -> None:
        caps = ModelCapabilities.load("deepseek_v4.flash", MODELS_DIR)
        # 2026-08-17 是周一；01:00 UTC 为峰时起点，04:00 已回到谷时。
        peak = resolve_chat_price(
            caps.pricing,
            occurred_at=datetime(2026, 8, 17, 1, 0, tzinfo=timezone.utc),
        )
        off_peak = resolve_chat_price(
            caps.pricing,
            occurred_at=datetime(2026, 8, 17, 4, 0, tzinfo=timezone.utc),
        )
        self.assertIsNotNone(peak)
        self.assertIsNotNone(off_peak)
        assert peak is not None and off_peak is not None
        self.assertEqual(peak.tier, "peak")
        self.assertEqual(peak.input_per_mtoken, Decimal("0.44"))
        self.assertEqual(off_peak.tier, "default")
        self.assertEqual(off_peak.output_per_mtoken, Decimal("0.66"))

    def test_deepseek_weekend_is_off_peak(self) -> None:
        caps = ModelCapabilities.load("deepseek_v4.pro", MODELS_DIR)
        quote = resolve_chat_price(
            caps.pricing,
            occurred_at=datetime(2026, 8, 22, 2, 0, tzinfo=timezone.utc),
        )
        assert quote is not None
        self.assertEqual(quote.tier, "default")

    def test_usd_quote_calculates_cache_split_in_cny(self) -> None:
        caps = ModelCapabilities.load("deepseek_v4.flash", MODELS_DIR)
        quote = resolve_chat_price(
            caps.pricing,
            occurred_at=datetime(2026, 8, 17, 1, 0, tzinfo=timezone.utc),
        )
        assert quote is not None
        result = calculate_chat_cost(
            quote,
            prompt_tokens=1_000_000,
            completion_tokens=500_000,
            cache_hit_tokens=800_000,
        )
        # miss:0.2*0.44*7.2 + hit:0.8*0.014*7.2 + out:0.5*1.32*7.2
        self.assertEqual(result.total_cny, Decimal("5.466240"))
        self.assertEqual(result.cache_miss_tokens, 200_000)

    def test_local_catalog_precedes_litellm_cost_map(self) -> None:
        caps = ModelCapabilities.load("deepseek_v4.flash", MODELS_DIR)
        added = []

        class FakeSession:
            def add(self, row):
                added.append(row)

        @contextmanager
        def fake_scope():
            yield FakeSession()

        with (
            patch("core.storage.usage.session_scope", fake_scope),
            patch("core.storage.usage._safe_chat_cost_usd", return_value=Decimal("999")),
        ):
            cost = record_chat_usage(
                task_id=None,
                user_id=uuid4(),
                message_id=None,
                model_profile="deepseek_v4.flash",
                prompt_tokens=1_000_000,
                completion_tokens=0,
                pricing=caps.pricing,
                occurred_at=datetime(2026, 8, 17, 4, 0, tzinfo=timezone.utc),
                response=SimpleNamespace(created=None),
            )

        self.assertEqual(cost, Decimal("1.584000"))
        self.assertEqual(added[0].units["pricing_revision"], "deepseek-v4-20260816")
        self.assertEqual(added[0].units["price_tier"], "default")


if __name__ == "__main__":
    unittest.main()