"""core/storage/usage_report.py 的 DB 级测试 —— 锁「计费读侧口径」。

背景:units JSONB 的写侧(usage.py)与读侧(usage_report.py)是同一契约的两半,
此前读侧散在 web 层硬编码、零测试,改 key 会静默算错。本测试用真实 PG 验证:
- cost 全 kind 合计、token/cache_hit 仅 chat
- task 维度批量聚合 / 按模型聚合 / 按用户分页聚合 三条读路径口径一致

无 DB(本机 PG 没起 / 未配)则整组 skip,不拖累纯单测环境。
数据纪律(公测期):只 INSERT 自己造的 user/task/usage_events,teardown 只
DELETE 这些行(按测试专属 user_id 过滤),绝不触碰既有数据。
"""
import os
import unittest
import uuid
from datetime import datetime, timedelta, timezone
from decimal import Decimal


def _test_db_ready() -> bool:
    """DB 级测试只认显式 `ZCBOT_TEST_DB_URL`,绝不回退 .env 的 ZCBOT_DB_URL。

    教训(2026-07-23 实锤):.env 的 ZCBOT_DB_URL 可能经隧道指向**生产库**(本次
    即 127.0.0.1:6012 → 生产 PG),测试若沿用它,插入的 enabled 到点 job 会被生产
    实例的调度守护认领并真跑 agent。测试库须显式指定:
        set ZCBOT_TEST_DB_URL=postgresql+psycopg://user:pass@127.0.0.1:5432/zcbot_test
    (建库后 `alembic upgrade head` 一次);未设 → 整组 skip。
    """
    url = os.environ.get("ZCBOT_TEST_DB_URL", "").strip()
    if not url:
        return False
    if "connect_timeout" not in url:
        # 测试库挂了要快速降级 skip,不能让整个 discovery 挂死在 TCP 建连上
        url += ("&" if "?" in url else "?") + "connect_timeout=5"
    os.environ["ZCBOT_DB_URL"] = url  # 本测试进程内覆盖,engine 单例随之指向测试库
    return True


try:
    if not _test_db_ready():
        raise RuntimeError("ZCBOT_TEST_DB_URL 未设")
    from core.storage import session_scope
    from core.storage.models import Task, UsageEvent, User
    from core.storage import usage_report

    with session_scope() as _s:
        _s.execute(__import__("sqlalchemy").select(1))
    _DB_OK = True
except Exception:
    _DB_OK = False


@unittest.skipUnless(_DB_OK, "ZCBOT_TEST_DB_URL 未设或测试库不可达,跳过 DB 级测试(绝不用 .env 的库)")
class UsageReportTests(unittest.TestCase):
    """一个测试专属 user + 两个 task,插一组已知 usage_events,验证三条读路径。"""

    @classmethod
    def setUpClass(cls):
        cls.uid = uuid.uuid4()
        cls.tid_a = uuid.uuid4()
        cls.tid_b = uuid.uuid4()
        with session_scope() as s:
            s.add(User(user_id=cls.uid, email=f"test-usage-report-{cls.uid.hex[:8]}@invalid.local"))
            s.flush()  # 无 relationship 映射,FK 依赖顺序要显式 flush 保证
            for tid, name in ((cls.tid_a, "ur-test-a"), (cls.tid_b, "ur-test-b")):
                s.add(Task(
                    task_id=tid, user_id=cls.uid, name=name,
                    working_dir=f"workspace/users/{cls.uid}/{name}",
                ))
            s.flush()
            # task A:两笔 chat(带缓存命中)+ 一笔 image(cost-only,token 不该被计入)
            s.add(UsageEvent(
                user_id=cls.uid, task_id=cls.tid_a, kind="chat",
                model_profile="ur-test.flash",
                units={"tokens_in": 1000, "tokens_out": 200, "cache_hit_tokens": 600},
                cost_cny=Decimal("0.010000"),
            ))
            s.add(UsageEvent(
                user_id=cls.uid, task_id=cls.tid_a, kind="chat",
                model_profile="ur-test.pro",
                units={"tokens_in": 3000, "tokens_out": 800, "cache_hit_tokens": 0},
                cost_cny=Decimal("0.200000"),
            ))
            s.add(UsageEvent(
                user_id=cls.uid, task_id=cls.tid_a, kind="image",
                model_profile="ur-test-seedream",
                units={"images": 1, "tokens_in": 999999},  # 非 chat 的 tokens 字段必须被忽略
                cost_cny=Decimal("0.300000"),
            ))
            # task B:一笔 chat;另一笔 task_id=NULL 的 kb_ingest(0022,不得进任何 task 聚合)
            s.add(UsageEvent(
                user_id=cls.uid, task_id=cls.tid_b, kind="chat",
                model_profile="ur-test.flash",
                units={"tokens_in": 500, "tokens_out": 100, "cache_hit_tokens": 250},
                cost_cny=Decimal("0.005000"),
            ))
            s.add(UsageEvent(
                user_id=cls.uid, task_id=None, kind="kb_ingest",
                model_profile="ur-test.flash",
                units={"kb": "测试库", "source": "x.pdf", "tokens_in": 42, "tokens_out": 7},
                cost_cny=Decimal("0.001000"),
            ))

    @classmethod
    def tearDownClass(cls):
        from sqlalchemy import delete
        with session_scope() as s:
            s.execute(delete(UsageEvent).where(UsageEvent.user_id == cls.uid))
            s.execute(delete(Task).where(Task.user_id == cls.uid))
            s.execute(delete(User).where(User.user_id == cls.uid))

    def test_task_aggregates_cost_all_kinds_tokens_chat_only(self):
        with session_scope() as s:
            agg = usage_report.task_usage_aggregates(s, [self.tid_a, self.tid_b])
        a = agg[self.tid_a]
        # cost = chat 0.01 + 0.20 + image 0.30;token 仅 chat(image 的 999999 忽略)
        self.assertAlmostEqual(a["cost_cny"], 0.51, places=6)
        self.assertEqual(a["tokens_in"], 4000)
        self.assertEqual(a["tokens_out"], 1000)
        self.assertEqual(a["tokens_cache_hit"], 600)
        b = agg[self.tid_b]
        self.assertAlmostEqual(b["cost_cny"], 0.005, places=6)
        self.assertEqual(b["tokens_in"], 500)
        self.assertEqual(b["tokens_cache_hit"], 250)

    def test_task_aggregates_empty_and_unknown(self):
        with session_scope() as s:
            self.assertEqual(usage_report.task_usage_aggregates(s, []), {})
            # 未知 task 无行 → 不出现在结果里(调用方 .get 兜 0)
            self.assertEqual(usage_report.task_usage_aggregates(s, [uuid.uuid4()]), {})

    def test_models_usage_groups_and_sorts(self):
        cutoff = datetime.now(timezone.utc) - timedelta(days=1)
        with session_scope() as s:
            rows = usage_report.models_usage(s, cutoff, sort="cost")
        by_mp = {r["model_profile"]: r for r in rows}
        # ur-test.flash:两笔 chat(task A 1000/200 + task B 500/100)+ 一笔 kb_ingest
        #(cost 计入、tokens 因非 chat 忽略)
        flash = by_mp["ur-test.flash"]
        self.assertEqual(flash["tokens_in"], 1500)
        self.assertEqual(flash["tokens_out"], 300)
        self.assertAlmostEqual(flash["cost_cny"], 0.016, places=6)
        self.assertEqual(flash["n_events"], 3)
        self.assertAlmostEqual(by_mp["ur-test.pro"]["cost_cny"], 0.2, places=6)
        self.assertAlmostEqual(by_mp["ur-test-seedream"]["cost_cny"], 0.3, places=6)
        self.assertEqual(by_mp["ur-test-seedream"]["tokens_in"], 0)

    def test_models_usage_cutoff_excludes_old(self):
        # cutoff 在未来 → 我们刚插的行全部被排除
        future = datetime.now(timezone.utc) + timedelta(days=1)
        with session_scope() as s:
            rows = usage_report.models_usage(s, future, sort="cost")
        self.assertNotIn("ur-test.flash", {r["model_profile"] for r in rows})

    def test_user_usage_page_finds_our_user(self):
        with session_scope() as s:
            d = usage_report.user_usage_page(s, page=0, page_size=100000, cutoff=None, sort="cost")
        mine = [r for r in d["rows"] if r["user_id"] == str(self.uid)]
        self.assertEqual(len(mine), 1)
        r = mine[0]
        # cost 全 kind:0.01+0.20+0.30+0.005+0.001;token 仅 chat:4500/1100;hit:850
        self.assertAlmostEqual(r["cost_cny"], 0.516, places=6)
        self.assertEqual(r["tokens_in"], 4500)
        self.assertEqual(r["tokens_out"], 1100)
        self.assertEqual(r["tokens_cache_hit"], 850)
        self.assertEqual(r["n_events"], 5)
        self.assertIsNotNone(r["last_used_at"])

    def test_usage_overview_delta(self):
        """overview 是全局聚合,共享 dev 库上断相对增量:排除我们行前后的差 = 我们插入的量。"""
        cutoff_7d = datetime.now(timezone.utc) - timedelta(days=7)
        with session_scope() as s:
            total = usage_report.usage_overview(s, cutoff_7d)["total"]
            # 我们贡献:cost 0.516 / tokens_in 4500 / tokens_out 1100 / hit 850 / 5 事件。
            # 断"至少包含"(库里还有真实数据,只验下界与口径不炸)。
            self.assertGreaterEqual(total["cost_cny"], 0.516 - 1e-6)
            self.assertGreaterEqual(total["tokens_in"], 4500)
            self.assertGreaterEqual(total["n_events"], 5)
            self.assertIsInstance(total["tokens_cache_hit"], int)


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