from decimal import Decimal
import unittest

from core.loop import _extract_usage_details
from core.storage.usage import _fallback_chat_cost_cny


class UsageAccountingTests(unittest.TestCase):
    def test_extract_usage_details_includes_cache_tokens(self) -> None:
        usage = {
            "prompt_tokens": 1200,
            "completion_tokens": 80,
            "prompt_cache_hit_tokens": 900,
            "prompt_cache_miss_tokens": 300,
            "completion_tokens_details": {"reasoning_tokens": 12},
        }

        details = _extract_usage_details(usage)

        self.assertEqual(details["tokens_in"], 1200)
        self.assertEqual(details["tokens_out"], 80)
        self.assertEqual(details["cache_hit_tokens"], 900)
        self.assertEqual(details["cache_miss_tokens"], 300)
        self.assertEqual(details["reasoning_tokens"], 12)

    def test_fallback_chat_cost_uses_price_snapshots(self) -> None:
        cost = _fallback_chat_cost_cny(
            prompt_tokens=1_000_000,
            completion_tokens=500_000,
            input_cny_per_mtoken=1.0,
            output_cny_per_mtoken=10.0,
        )

        self.assertEqual(cost, Decimal("6.000000"))

    def test_fallback_chat_cost_discounts_cache_hits(self) -> None:
        # 100 万输入里 80 万命中缓存(0.1 价),20 万未命中(1.0 价),50 万输出(10 价)
        cost = _fallback_chat_cost_cny(
            prompt_tokens=1_000_000,
            completion_tokens=500_000,
            input_cny_per_mtoken=1.0,
            output_cny_per_mtoken=10.0,
            cache_hit_tokens=800_000,
            cache_hit_cny_per_mtoken=0.1,
        )
        # 0.2(miss) + 0.08(hit) + 5.0(out) = 5.28
        self.assertEqual(cost, Decimal("5.280000"))

    def test_fallback_chat_cost_no_cache_price_charges_full(self) -> None:
        # 未配缓存价(0)→ 命中段不打折,按 input 全价(老行为,绝不少记)
        cost = _fallback_chat_cost_cny(
            prompt_tokens=1_000_000,
            completion_tokens=0,
            input_cny_per_mtoken=1.0,
            output_cny_per_mtoken=10.0,
            cache_hit_tokens=900_000,
            cache_hit_cny_per_mtoken=0.0,
        )
        self.assertEqual(cost, Decimal("1.000000"))


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