"""协调员子任务创建时配额门单测。"""
import json
import os
import sqlite3
import tempfile
import unittest

from cloud.server import init_db
from cloud.worker import Worker


def _insert_parent(db_path, task_id="coord", created_by="owner1"):
    conn = sqlite3.connect(db_path)
    try:
        parts = json.dumps([{"name": "后端", "prompt": "做后端"},
                            {"name": "前端", "prompt": "做前端"}])
        conn.execute(
            "INSERT INTO tasks(id, project_id, prompt, status,"
            " created_by, created_at, updated_at, coordinator_parts)"
            " VALUES(?, 'proj', 'coord', 'running', ?, ?, ?, ?)",
            (task_id, created_by, 1000.0, 1000.0, parts))
        conn.commit()
    finally:
        conn.close()


def _children(db_path, parent_id):
    conn = sqlite3.connect(db_path)
    try:
        return conn.execute(
            "SELECT id FROM tasks WHERE parent_task_id = ?",
            (parent_id,)).fetchall()
    finally:
        conn.close()


def _audit_count(db_path, action):
    conn = sqlite3.connect(db_path)
    try:
        try:
            return conn.execute(
                "SELECT COUNT(*) FROM audit_log WHERE action = ?",
                (action,)).fetchone()[0]
        except sqlite3.OperationalError:
            # 审计是 best-effort：写审计的连接与任务写连接在
            # SQLite 上撞写锁时审计会静默丢失（既有行为），
            # 此时表可能根本不存在，计 0 而不是让测试炸掉。
            return 0
    finally:
        conn.close()


def _set_quota(db_path, user_id, micros):
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "CREATE TABLE IF NOT EXISTS user_quotas("
            "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
            " updated_at DOUBLE PRECISION)")
        conn.execute(
            "INSERT OR REPLACE INTO user_quotas(user_id,"
            " quota_usd_micros, updated_at) VALUES(?, ?, ?)",
            (user_id, micros, 1000.0))
        conn.commit()
    finally:
        conn.close()


class CoordQuotaGateTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        init_db(self.db_path)
        self.env = {"XCM_API_KEY": "test-placeholder"}

    def tearDown(self):
        self.tmp.cleanup()

    def _worker(self):
        return Worker(self.db_path, "/fake/host.py", env=self.env)

    def _parts(self):
        return [{"name": "后端", "prompt": "做后端"},
                {"name": "前端", "prompt": "做前端"}]

    def test_skip_when_over_quota(self):
        _insert_parent(self.db_path)
        _set_quota(self.db_path, "owner1", 0)
        worker = self._worker()
        task = worker.get_task("coord")
        children = worker._coord_sync_children(
            task, task["coordinator_parts"])
        self.assertEqual(children, [])
        self.assertEqual(_children(self.db_path, "coord"), [])
        # 两个分工各记一次跳过审计；第二轮重调不再重复记。
        self.assertEqual(
            _audit_count(self.db_path, "coordinator_child_quota_skipped"),
            2)
        worker._coord_sync_children(task, task["coordinator_parts"])
        self.assertEqual(
            _audit_count(self.db_path, "coordinator_child_quota_skipped"),
            2)
        self.assertEqual(_children(self.db_path, "coord"), [])

    def test_create_when_within_quota(self):
        _insert_parent(self.db_path)
        _set_quota(self.db_path, "owner1", 100000000)
        worker = self._worker()
        task = worker.get_task("coord")
        children = worker._coord_sync_children(
            task, task["coordinator_parts"])
        self.assertEqual(len(children), 2)
        self.assertEqual(len(_children(self.db_path, "coord")), 2)
        self.assertEqual(
            _audit_count(self.db_path, "coordinator_child_quota_skipped"),
            0)

    def test_create_with_default_quota_when_no_row(self):
        # 无配额行按默认 $10 计，未花过即放行（不是不限额，
        # 也不是零额度误拒）。
        _insert_parent(self.db_path)
        worker = self._worker()
        task = worker.get_task("coord")
        children = worker._coord_sync_children(
            task, task["coordinator_parts"])
        self.assertEqual(len(children), 2)


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