"""并发上限单测（caps）。

覆盖：配置解析优先级与默认值、认领门的未满 / 已满 / 释放后
恢复、启动中任务计入占用、CAS 抢输后占位不泄漏、缺 key 门
提前返回后占位释放、双线程同时认领只赢一个。

只用注入的假环境与本地 SQLite，不跑真实模型、不联网。
"""

import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path

try:
    from cloud.caps import (
        can_claim,
        effective_capacity,
        resolve_capacity,
        running_count,
    )
    from cloud.server import init_db
    from cloud.worker import Worker
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.caps import (
        can_claim,
        effective_capacity,
        resolve_capacity,
        running_count,
    )
    from cloud.server import init_db
    from cloud.worker import Worker


def _insert_task(db_path, task_id, project_id="p1"):
    now = time.time()
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO tasks(id, project_id, prompt, status,"
            " idempotency_key, depends_on, priority, created_by,"
            " created_at, updated_at, last_heartbeat)"
            " VALUES(?, ?, ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)",
            (task_id, project_id, "do work", now, now),
        )
        conn.commit()
    finally:
        conn.close()


def _status(db_path, task_id):
    conn = sqlite3.connect(db_path)
    try:
        row = conn.execute(
            "SELECT status FROM tasks WHERE id = ?",
            (task_id,)).fetchone()
    finally:
        conn.close()
    return row[0] if row else None


class CapsResolveTestCase(unittest.TestCase):
    def test_explicit_beats_env(self):
        env = {"CLOUD_WORKER_CAPACITY": "3"}
        self.assertEqual(resolve_capacity(2, env=env), 2)

    def test_env_used_when_no_explicit(self):
        env = {"CLOUD_WORKER_CAPACITY": "3"}
        self.assertEqual(resolve_capacity(None, env=env), 3)

    def test_default_is_1(self):
        self.assertEqual(resolve_capacity(None, env={}), 1)
        self.assertEqual(resolve_capacity(), 1)

    def test_invalid_values_fall_back_to_1(self):
        self.assertEqual(resolve_capacity(0, env={}), 1)
        self.assertEqual(resolve_capacity(-2, env={}), 1)
        self.assertEqual(resolve_capacity("abc", env={}), 1)
        self.assertEqual(resolve_capacity(None, env={
            "CLOUD_WORKER_CAPACITY": "abc"}), 1)
        self.assertEqual(resolve_capacity(None, env={
            "CLOUD_WORKER_CAPACITY": "0"}), 1)
        self.assertEqual(resolve_capacity(None, env={
            "CLOUD_WORKER_CAPACITY": "  "}), 1)

    def test_effective_capacity(self):
        self.assertEqual(effective_capacity(2), 2)
        self.assertEqual(effective_capacity(0), 1)
        self.assertEqual(effective_capacity(None), 1)

    def test_can_claim_boundaries(self):
        self.assertTrue(can_claim(0, 1))
        self.assertFalse(can_claim(1, 1))
        self.assertTrue(can_claim(1, 2))
        self.assertFalse(can_claim(2, 2))

    def test_running_count_counts_starting_entries(self):
        # 空条目 = 已认领未起进程的启动中任务，同样占一个槽。
        self.assertEqual(running_count({"a": {}, "b": {"proc": 1}}), 2)
        self.assertEqual(running_count({}), 0)


class WorkerClaimGateTestCase(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, **kwargs):
        kwargs.setdefault("env", self.env)
        kwargs.setdefault("retry_backoff_seconds", 0)
        kwargs.setdefault("max_retries", 0)
        return Worker(self.db_path, "/fake/host.py", **kwargs)

    def test_default_capacity_is_1(self):
        worker = self._worker()
        self.assertEqual(worker.capacity, 1)

    def test_env_capacity_and_explicit_priority(self):
        env = {"CLOUD_WORKER_CAPACITY": "3"}
        self.assertEqual(self._worker(env=env).capacity, 3)
        self.assertEqual(self._worker(env=env, capacity=2).capacity, 2)

    def test_full_worker_refuses_claim_and_task_stays_queued(self):
        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        worker = self._worker(capacity=1)
        self.assertTrue(worker._claim("A"))
        self.assertEqual(_status(self.db_path, "A"), "running")
        self.assertFalse(worker._claim("B"))
        self.assertEqual(_status(self.db_path, "B"), "queued")
        self.assertEqual(set(worker._running_snapshot()), {"A"})

    def test_release_restores_claim(self):
        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        worker = self._worker(capacity=1)
        self.assertTrue(worker._claim("A"))
        self.assertFalse(worker._claim("B"))
        worker._running_clear("A")
        self.assertTrue(worker._claim("B"))
        self.assertEqual(_status(self.db_path, "B"), "running")

    def test_starting_task_counts_toward_capacity(self):
        _insert_task(self.db_path, "B")
        worker = self._worker(capacity=1)
        # 空条目模拟“已认领、执行器还没起进程”的启动中任务。
        worker._running_set("STARTING")
        self.assertEqual(worker._running_count(), 1)
        self.assertFalse(worker._claim("B"))
        self.assertEqual(_status(self.db_path, "B"), "queued")
        worker._running_clear("STARTING")
        self.assertTrue(worker._claim("B"))

    def test_capacity_2_allows_two_blocks_third(self):
        for task_id in ("A", "B", "C"):
            _insert_task(self.db_path, task_id)
        worker = self._worker(capacity=2)
        self.assertTrue(worker._claim("A"))
        self.assertTrue(worker._claim("B"))
        self.assertFalse(worker._claim("C"))
        self.assertEqual(_status(self.db_path, "C"), "queued")
        worker._running_clear("A")
        self.assertTrue(worker._claim("C"))

    def test_lost_cas_race_releases_slot(self):
        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        first = self._worker(capacity=1)
        second = self._worker(capacity=1)
        self.assertTrue(first._claim("A"))
        # second 的 CAS 抢输：返回 False，且它自己的占位已清，
        # 不能留下一个空槽把自己堵死。
        self.assertFalse(second._claim("A"))
        self.assertEqual(second._running_snapshot(), {})
        self.assertTrue(second._claim("B"))

    def test_two_threads_claim_only_one_wins(self):
        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        worker = self._worker(capacity=1)
        barrier = threading.Barrier(2, timeout=5)
        results = {}

        def _try(task_id):
            barrier.wait()
            results[task_id] = worker._claim(task_id)

        threads = [threading.Thread(target=_try, args=(t,))
                   for t in ("A", "B")]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=10)
        self.assertEqual(sorted(results.values()), [False, True])
        self.assertEqual(worker._running_count(), 1)

    def test_missing_key_gate_releases_slot(self):
        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        # env 里无 key：run_one 在缺 key 门提前返回 failed，
        # 占位必须随之释放，否则 capacity=1 的 worker 会被堵死。
        worker = self._worker(env={}, capacity=1)
        result = worker.run_one()
        self.assertTrue(result["dispatched"])
        self.assertEqual(result["status"], "failed")
        self.assertEqual(_status(self.db_path, "A"), "failed")
        self.assertEqual(worker._running_snapshot(), {})
        self.assertTrue(worker._claim("B"))
        self.assertEqual(_status(self.db_path, "B"), "running")


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