"""租户配额与计费单测：月度 USD 上限 + 准入 402 + 管理端点。

不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
SQLite；用量直接插 cost_records 行（按任务 created_by 归属、
UTC day），验证 JOIN 归属、自然月窗口与三路准入（建任务 /
分支 / 重试）。配额单位是微美元（1 美元 = 1_000_000），
与 cost_records 一致。
"""

import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
from datetime import date, datetime, timedelta, timezone
from pathlib import Path

try:
    from cloud.auth import AuthStore
    from cloud.server import create_server
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.server import create_server


def _req(port, method, path, body=None, key=None):
    conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10)
    hdrs = {}
    payload = None
    if body is not None:
        payload = json.dumps(body).encode("utf-8")
        hdrs["Content-Type"] = "application/json"
    if key is not None:
        hdrs["X-API-Key"] = key
    try:
        conn.request(method, path, body=payload, headers=hdrs)
        resp = conn.getresponse()
        data = resp.read()
    finally:
        conn.close()
    try:
        parsed = json.loads(data.decode("utf-8")) if data else None
    except ValueError:
        parsed = None
    return resp.status, parsed


def _set_task_status(db_path, task_id, status):
    conn = sqlite3.connect(db_path)
    try:
        conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
                     (status, task_id))
        conn.commit()
    finally:
        conn.close()


def _insert_cost(db_path, task_id, project_id, usd_micros, day):
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO cost_records(task_id, project_id, model,"
            " input_tokens, output_tokens, total_tokens, usd_micros,"
            " source, day, created_at) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (task_id, project_id, "gpt-5.6-terra", 100, 50, 150,
             usd_micros, "actual", day, 0))
        conn.commit()
    finally:
        conn.close()


def _utc_today():
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def _prev_month_day():
    today = date.today()
    first = today.replace(day=1)
    return (first - timedelta(days=1)).isoformat()


class ServerQuotasTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        self.auth = AuthStore(self.db_path)
        self.owner_id, self.owner_key = self.auth.create_user("owner")
        self.other_id, self.other_key = self.auth.create_user("other")
        self.server = create_server(self.db_path, "127.0.0.1", 0)
        self.port = self.server.server_address[1]
        self.thread = threading.Thread(
            target=self.server.serve_forever, daemon=True)
        self.thread.start()
        s, proj = _req(self.port, "POST", "/projects",
                       {"name": "qproj"}, self.owner_key)
        self.assertEqual(s, 201, proj)
        self.project_id = proj["id"]

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.tmp.cleanup()

    def _create_task(self, key=None):
        return _req(self.port, "POST", "/tasks",
                    {"project_id": self.project_id,
                     "prompt": "quota probe"},
                    key if key is not None else self.owner_key)

    def _set_quota(self, target, quota, key=None):
        return _req(self.port, "POST", "/admin/quotas",
                    {"user_id": target, "quota_usd_micros": quota},
                    key if key is not None else self.owner_key)

    def test_self_quota_default_unlimited(self):
        s, q = _req(self.port, "GET", "/users/me/quota",
                    key=self.owner_key)
        self.assertEqual(s, 200, q)
        self.assertIsNone(q["quota_usd_micros"])
        self.assertEqual(q["used_usd_micros"], 0)
        self.assertIsNone(q["remaining_usd_micros"])
        self.assertEqual(q["month"],
                         datetime.now(timezone.utc).strftime("%Y-%m"))

    def test_admin_set_and_get(self):
        s, q = self._set_quota(self.owner_id, 5000000)
        self.assertEqual(s, 200, q)
        self.assertEqual(q["quota_usd_micros"], 5000000)
        s, q = _req(self.port, "GET",
                    "/admin/quotas?user_id=" + self.owner_id,
                    key=self.owner_key)
        self.assertEqual(s, 200, q)
        self.assertEqual(q["quota_usd_micros"], 5000000)
        s, q = _req(self.port, "GET", "/users/me/quota",
                    key=self.owner_key)
        self.assertEqual(q["quota_usd_micros"], 5000000)
        # null = 恢复不限额
        s, q = self._set_quota(self.owner_id, None)
        self.assertEqual(s, 200, q)
        self.assertIsNone(q["quota_usd_micros"])

    def test_admin_set_forbidden_and_validation(self):
        # other 没有任何项目角色，不是管理员
        s, q = self._set_quota(self.owner_id, 100, self.other_key)
        self.assertEqual(s, 403, q)
        s, q = _req(self.port, "GET",
                    "/admin/quotas?user_id=" + self.owner_id,
                    key=self.other_key)
        self.assertEqual(s, 403, q)
        s, q = self._set_quota(self.owner_id, -1)
        self.assertEqual(s, 400, q)
        s, q = self._set_quota("no-such-user", 100)
        self.assertEqual(s, 404, q)
        s, q = _req(self.port, "POST", "/admin/quotas",
                    {"quota_usd_micros": 100}, self.owner_key)
        self.assertEqual(s, 400, q)
        s, q = _req(self.port, "GET", "/admin/quotas",
                    key=self.owner_key)
        self.assertEqual(s, 400, q)

    def test_usage_counts_current_month_only(self):
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        _insert_cost(self.db_path, task["id"], self.project_id,
                     5000000, _utc_today())
        _insert_cost(self.db_path, task["id"], self.project_id,
                     100000000, _prev_month_day())
        s, q = _req(self.port, "GET", "/users/me/quota",
                    key=self.owner_key)
        self.assertEqual(s, 200, q)
        self.assertEqual(q["used_usd_micros"], 5000000)

    def test_create_blocked_at_quota(self):
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        _insert_cost(self.db_path, task["id"], self.project_id,
                     5000000, _utc_today())
        s, _ = self._set_quota(self.owner_id, 5000000)
        self.assertEqual(s, 200)
        s, body = self._create_task()
        self.assertEqual(s, 402, body)
        self.assertEqual(body["error"], "quota_exceeded")
        self.assertEqual(body["quota_usd_micros"], 5000000)
        self.assertEqual(body["used_usd_micros"], 5000000)
        # 额度调高后放行
        s, _ = self._set_quota(self.owner_id, 6000000)
        self.assertEqual(s, 200)
        s, task2 = self._create_task()
        self.assertEqual(s, 201, task2)

    def test_fork_blocked_at_quota(self):
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        _insert_cost(self.db_path, task["id"], self.project_id,
                     3000000, _utc_today())
        _set_task_status(self.db_path, task["id"], "completed")
        s, _ = self._set_quota(self.owner_id, 3000000)
        self.assertEqual(s, 200)
        s, body = _req(self.port, "POST",
                       "/tasks/%s/fork" % task["id"],
                       {"message": "branch probe"}, self.owner_key)
        self.assertEqual(s, 402, body)
        s, _ = self._set_quota(self.owner_id, None)
        self.assertEqual(s, 200)
        s, forked = _req(self.port, "POST",
                         "/tasks/%s/fork" % task["id"],
                         {"message": "branch probe"}, self.owner_key)
        self.assertEqual(s, 201, forked)

    def test_retry_blocked_at_quota(self):
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        _insert_cost(self.db_path, task["id"], self.project_id,
                     2000000, _utc_today())
        _set_task_status(self.db_path, task["id"], "failed")
        s, _ = self._set_quota(self.owner_id, 2000000)
        self.assertEqual(s, 200)
        s, body = _req(self.port, "POST",
                       "/tasks/%s/retry" % task["id"], {},
                       self.owner_key)
        self.assertEqual(s, 402, body)
        s, row = _req(self.port, "GET", "/tasks/" + task["id"],
                      key=self.owner_key)
        self.assertEqual(row["status"], "failed")
        s, _ = self._set_quota(self.owner_id, None)
        self.assertEqual(s, 200)
        s, retried = _req(self.port, "POST",
                          "/tasks/%s/retry" % task["id"], {},
                          self.owner_key)
        self.assertEqual(s, 200, retried)
        self.assertEqual(retried["status"], "queued")

    def test_retry_uses_creator_quota_not_operator(self):
        # 重试按任务创建者查配额：费用记创建者名下，操作人
        # 是谁不应影响准入。other 须是 maintainer 才能过
        # approve_merge 权限门，否则 403 会掩盖配额判断。
        self.auth.set_member(self.project_id, self.owner_id,
                             self.other_id, "maintainer")
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        _insert_cost(self.db_path, task["id"], self.project_id,
                     2000000, _utc_today())
        _set_task_status(self.db_path, task["id"], "failed")
        # 创建者超限：维护者替他重试也应拦。
        s, _ = self._set_quota(self.owner_id, 2000000)
        self.assertEqual(s, 200)
        s, body = _req(self.port, "POST",
                       "/tasks/%s/retry" % task["id"], {},
                       self.other_key)
        self.assertEqual(s, 402, body)
        # 创建者不限额、操作人额度为 0：按旧口径（查操作人）
        # 会拦，按创建者口径应放行。
        s, _ = self._set_quota(self.owner_id, None)
        self.assertEqual(s, 200)
        s, _ = self._set_quota(self.other_id, 0)
        self.assertEqual(s, 200)
        s, retried = _req(self.port, "POST",
                          "/tasks/%s/retry" % task["id"], {},
                          self.other_key)
        self.assertEqual(s, 200, retried)
        self.assertEqual(retried["status"], "queued")

    def test_continue_blocked_at_creator_quota(self):
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        _insert_cost(self.db_path, task["id"], self.project_id,
                     5000000, _utc_today())
        _set_task_status(self.db_path, task["id"], "completed")
        s, _ = self._set_quota(self.owner_id, 5000000)
        self.assertEqual(s, 200)
        # 以前：先回 200 排队、随后认领门悄悄标 interrupted。
        # 现在：API 侧直接 402，状态保持 completed 不变。
        s, body = _req(self.port, "POST",
                       "/tasks/%s/continue" % task["id"],
                       {"message": "again"}, self.owner_key)
        self.assertEqual(s, 402, body)
        s, row = _req(self.port, "GET", "/tasks/" + task["id"],
                      key=self.owner_key)
        self.assertEqual(row["status"], "completed")
        # 额度调高后放行。
        s, _ = self._set_quota(self.owner_id, None)
        self.assertEqual(s, 200)
        s, continued = _req(self.port, "POST",
                            "/tasks/%s/continue" % task["id"],
                            {"message": "again"}, self.owner_key)
        self.assertEqual(s, 200, continued)
        self.assertEqual(continued["status"], "queued")
        # 续跑同样按创建者查：操作人额度为 0、创建者不限额时放行。
        self.auth.set_member(self.project_id, self.owner_id,
                             self.other_id, "maintainer")
        s, task2 = self._create_task()
        self.assertEqual(s, 201, task2)
        _set_task_status(self.db_path, task2["id"], "completed")
        s, _ = self._set_quota(self.other_id, 0)
        self.assertEqual(s, 200)
        s, continued2 = _req(self.port, "POST",
                             "/tasks/%s/continue" % task2["id"],
                             {"message": "again"}, self.other_key)
        self.assertEqual(s, 200, continued2)
        self.assertEqual(continued2["status"], "queued")

    def test_quota_reports_unpriced_tokens(self):
        s, task = self._create_task()
        self.assertEqual(s, 201, task)
        # 有价 5000000 计入 used；无价 actual 150 只进
        # unpriced_actual_tokens，不计价、不拦任务。
        _insert_cost(self.db_path, task["id"], self.project_id,
                     5000000, _utc_today())
        _insert_cost(self.db_path, task["id"], self.project_id,
                     None, _utc_today())
        s, q = _req(self.port, "GET", "/users/me/quota",
                    key=self.owner_key)
        self.assertEqual(s, 200, q)
        self.assertEqual(q["used_usd_micros"], 5000000)
        self.assertEqual(q["unpriced_actual_tokens"], 150)
        s, q = _req(self.port, "GET",
                    "/admin/quotas?user_id=" + self.owner_id,
                    key=self.owner_key)
        self.assertEqual(s, 200, q)
        self.assertEqual(q["unpriced_actual_tokens"], 150)


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