"""多 worker 池派发 + 通用 CloudExecutor + 自动 PR 单测。

只用标准库，不跑真实模型回合、不联网、不登录任何服务；
git/gh 一律用假 runner 注入，不依赖本机是否装 gh、是否有
远端仓库；CloudExecutor 的 endpoint 一律用假 http_client。
"""

import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

try:
    from cloud.auth import AuthStore
    from cloud.pr import PrSkipped, create_task_pr
    from cloud.sandbox import CloudExecutor, DockerExecutor
    from cloud.server import create_server, init_db
    from cloud.worker import Worker
    from cloud.workers import (
        heartbeat_worker,
        list_workers,
        register_worker,
    )
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.pr import PrSkipped, create_task_pr
    from cloud.sandbox import CloudExecutor, DockerExecutor
    from cloud.server import create_server, init_db
    from cloud.worker import Worker
    from cloud.workers import (
        heartbeat_worker,
        list_workers,
        register_worker,
    )


def _result(returncode=0, stdout="", stderr=""):
    return SimpleNamespace(returncode=returncode, stdout=stdout,
                           stderr=stderr)


class _ServerCase(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.viewer_id, self.viewer_key = self.auth.create_user("viewer")
        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()
        status, self.project = self._req(
            "POST", "/projects", {"name": "p1", "auto_pr": 1},
            self.owner_key)
        self.assertEqual(status, 201, self.project)
        self.auth.set_member(
            self.project["id"], self.owner_id, self.viewer_id, "viewer")
        self.worktrees = os.path.join(self.tmp.name, "worktrees")
        os.makedirs(self.worktrees, exist_ok=True)
        self.server.host_worktrees_dir = self.worktrees

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)
        try:
            self.auth.conn.close()
        except Exception:
            pass
        self.tmp.cleanup()

    def _req(self, method, path, body=None, key=None, headers=None):
        conn = http.client.HTTPConnection(
            "127.0.0.1", self.port, timeout=10)
        hdrs = dict(headers or {})
        payload = None
        if body is not None:
            payload = json.dumps(body).encode("utf-8")
            hdrs.setdefault("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 _insert_task(self, task_id, status, project_id=None):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "INSERT INTO tasks(id, project_id, prompt, status,"
                " depends_on, priority, created_by, created_at,"
                " updated_at, last_heartbeat) VALUES(?, ?, ?, ?, '[]',"
                " 0, 'system', 0, 0, NULL)",
                (task_id, project_id or self.project["id"],
                 "do it", status))
            conn.commit()
        finally:
            conn.close()

    def _task_row(self, task_id):
        conn = sqlite3.connect(self.db_path)
        try:
            row = conn.execute(
                "SELECT status, pr_url, pr_number FROM tasks"
                " WHERE id = ?", (task_id,)).fetchone()
        finally:
            conn.close()
        return row

    def _audit_actions(self):
        conn = sqlite3.connect(self.db_path)
        try:
            return [r[0] for r in conn.execute(
                "SELECT action FROM audit_log").fetchall()]
        finally:
            conn.close()

    def _git_runner(self, push_fail=False, remote_fail=False):
        def run(cmd, cwd):
            if "remote" in cmd:
                if remote_fail:
                    return _result(1, "", "no remote")
                return _result(0, "https://example.com/repo.git\n")
            if "branch" in cmd:
                # 假 worktree 正好在自己的 host/<任务名> 分支上：
                # worktree 目录名就是任务 id（见 _worktree）。
                return _result(0, "host/%s\n" % os.path.basename(cwd or ""))
            if "add" in cmd or "commit" in cmd:
                return _result(0, "")
            if "push" in cmd:
                if push_fail:
                    return _result(1, "", "push rejected")
                return _result(0, "")
            return _result(0, "")
        return run

    def _gh_runner(self, authed=True, pr_url=None):
        def run(cmd, cwd):
            if cmd[:2] == ["gh", "auth"]:
                return _result(0 if authed else 1, "", "")
            if cmd[:2] == ["gh", "pr"]:
                return _result(0, (pr_url or
                                   "https://example.com/repo/pull/42\n"))
            return _result(0, "")
        return run

    def _worktree(self, task_id):
        path = os.path.join(self.worktrees, task_id)
        os.makedirs(path, exist_ok=True)
        return path


class WorkerRegistryTestCase(_ServerCase):
    def test_heartbeat_requires_worker_id(self):
        status, body = self._req(
            "POST", "/workers/heartbeat", {}, self.owner_key)
        self.assertEqual(status, 400, body)

    def test_heartbeat_registers_and_lists_alive(self):
        status, body = self._req(
            "POST", "/workers/heartbeat",
            {"worker_id": "w-docker-1", "executor": "docker",
             "capacity": 2},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertTrue(body["alive"])
        self.assertEqual(body["executor"], "docker")
        status, body = self._req("GET", "/workers", None, self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(len(body["workers"]), 1)
        self.assertTrue(body["workers"][0]["alive"])

    def test_heartbeat_preserves_executor_when_not_given(self):
        self._req("POST", "/workers/heartbeat",
                  {"worker_id": "w1", "executor": "docker"}, self.owner_key)
        status, body = self._req(
            "POST", "/workers/heartbeat", {"worker_id": "w1"},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["executor"], "docker")

    def test_stale_worker_lists_not_alive(self):
        register_worker(self.db_path, "w-old", executor="local")
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "UPDATE workers SET last_seen = ? WHERE worker_id = ?",
                (time.time() - 3600, "w-old"))
            conn.commit()
        finally:
            conn.close()
        workers = list_workers(self.db_path)
        self.assertEqual(len(workers), 1)
        self.assertFalse(workers[0]["alive"])

    def test_heartbeat_without_key_is_unauthorized(self):
        status, body = self._req(
            "POST", "/workers/heartbeat", {"worker_id": "w1"}, None)
        self.assertEqual(status, 401, body)


class TaskExecutorFilterTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        init_db(os.path.join(self.tmp.name, "control.db"))
        self.db_path = os.path.join(self.tmp.name, "control.db")
        conn = sqlite3.connect(self.db_path)
        try:
            for tid, exc in (("t-local", "local"), ("t-none", None),
                             ("t-docker", "docker")):
                conn.execute(
                    "INSERT INTO tasks(id, project_id, prompt, status,"
                    " depends_on, priority, created_by, created_at,"
                    " updated_at, last_heartbeat, executor)"
                    " VALUES(?, 'proj', 'do it', 'queued', '[]', 0,"
                    " 'system', 0, 0, NULL, ?)", (tid, exc))
            conn.commit()
        finally:
            conn.close()

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

    def _docker_worker(self, capacity=None):
        return Worker(
            self.db_path, "host.py", config_path=None, repo=None,
            executor=DockerExecutor(runner=lambda cmd: _result()),
            env={"XCM_API_KEY": "test-placeholder"}, env_key=None,
            task_timeout=1800, max_retries=0, retry_backoff_seconds=0,
            capacity=capacity)

    def test_claim_rejects_other_pool_task(self):
        worker = self._docker_worker()
        self.assertFalse(worker._claim("t-local", executor="docker"))
        conn = sqlite3.connect(self.db_path)
        try:
            status = conn.execute(
                "SELECT status FROM tasks WHERE id = 't-local'"
            ).fetchone()[0]
        finally:
            conn.close()
        self.assertEqual(status, "queued")

    def test_claim_accepts_unassigned_and_same_pool(self):
        # 本测试连续认领两个任务, 测的是执行器过滤不是上限;
        # caps 默认上限 1 会挡第二单, 故显式给 capacity=2。
        worker = self._docker_worker(capacity=2)
        self.assertTrue(worker._claim("t-none", executor="docker"))
        self.assertTrue(worker._claim("t-docker", executor="docker"))

    def test_cas_race_only_one_worker_wins(self):
        w1 = self._docker_worker()
        w2 = self._docker_worker()
        self.assertTrue(w1._claim("t-none", executor="docker"))
        self.assertFalse(w2._claim("t-none", executor="docker"))

    def test_pick_next_skips_other_pool(self):
        # 只剩 local 池任务时，docker worker 无就绪任务。
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute("DELETE FROM tasks WHERE id != 't-local'")
            conn.commit()
        finally:
            conn.close()
        worker = self._docker_worker()
        self.assertIsNone(worker.pick_next())

    def test_run_one_skips_other_pool_without_model(self):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute("DELETE FROM tasks WHERE id != 't-local'")
            conn.commit()
        finally:
            conn.close()
        worker = self._docker_worker()
        result = worker.run_one()
        self.assertFalse(result["dispatched"])
        self.assertEqual(result["reason"], "no_ready_task")


class CloudExecutorTestCase(unittest.TestCase):
    def test_unconfigured_fails_without_fallback(self):
        ex = CloudExecutor(endpoint=None, env={})
        with self.assertRaises(RuntimeError) as ctx:
            ex.execute({"id": "t1", "prompt": "do it"})
        self.assertIn("not configured", str(ctx.exception))

    def test_configured_posts_only_id_prompt_repo_no_secrets(self):
        seen = {}

        def fake_client(url, payload, headers):
            seen["url"] = url
            seen["payload"] = payload
            seen["headers"] = headers
            return {"status": "completed", "stdout": "done"}

        ex = CloudExecutor(
            endpoint="https://runner.example.invalid/run",
            http_client=fake_client,
            auth_header="X-Runner-Token",
            env={"CLOUD_RUNNER_AUTH_TOKEN": "test-token-value"})
        result = ex.execute({"id": "t9", "prompt": "do it",
                             "repo": "https://example.com/r.git"})
        self.assertEqual(result["status"], "completed")
        self.assertEqual(seen["url"], "https://runner.example.invalid/run")
        self.assertEqual(set(seen["payload"].keys()),
                         {"id", "prompt", "repo_url"})
        self.assertEqual(seen["payload"]["repo_url"],
                         "https://example.com/r.git")
        self.assertNotIn("test-token-value",
                         json.dumps(seen["payload"]))
        self.assertEqual(seen["headers"]["X-Runner-Token"],
                         "test-token-value")
        # dispatch_task 是 worker 走的入口，必须与 execute 同结果。
        result2 = ex.dispatch_task({"id": "t9", "prompt": "do it"})
        self.assertEqual(result2["status"], "completed")

    def test_failed_status_maps_to_failed(self):
        ex = CloudExecutor(
            endpoint="https://runner.example.invalid/run",
            http_client=lambda u, p, h: {"status": "failed",
                                          "stdout": "boom"})
        result = ex.execute({"id": "t1", "prompt": "x"})
        self.assertEqual(result["status"], "failed")


class AutoPrTestCase(_ServerCase):
    def _pr(self, task_id, key):
        return self._req("POST", "/tasks/%s/pr" % task_id, {},
                         key)

    def test_auto_pr_disabled_skips_and_stays_merged(self):
        status, proj = self._req(
            "POST", "/projects", {"name": "p-off"}, self.owner_key)
        self.assertEqual(status, 201, proj)
        self._insert_task("toff", "merged", project_id=proj["id"])
        status, body = self._pr("toff", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertTrue(body["pr_skipped"])
        self.assertEqual(body["reason"], "auto_pr_disabled")
        self.assertEqual(self._task_row("toff")[0], "merged")

    def test_auto_pr_illegal_value_rejected(self):
        status, body = self._req(
            "POST", "/projects", {"name": "p-bad", "auto_pr": 2},
            self.owner_key)
        self.assertEqual(status, 400, body)

    def test_non_merged_returns_409(self):
        self._insert_task("tq", "queued")
        status, body = self._pr("tq", self.owner_key)
        self.assertEqual(status, 409, body)

    def test_viewer_cannot_pr(self):
        self._insert_task("tm", "merged")
        status, body = self._pr("tm", self.viewer_key)
        self.assertEqual(status, 403, body)

    def test_no_worktree_skips(self):
        self._insert_task("tnw", "merged")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tnw", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "no_worktree")
        self.assertEqual(self._task_row("tnw")[0], "merged")

    def test_no_remote_skips(self):
        self._insert_task("tnr", "merged")
        self._worktree("tnr")
        self.server.pr_git_runner = self._git_runner(remote_fail=True)
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tnr", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "no_remote")

    def test_no_gh_skips(self):
        self._insert_task("tngh", "merged")
        self._worktree("tngh")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = ""
        status, body = self._pr("tngh", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "no_gh")

    def test_gh_not_authenticated_skips(self):
        self._insert_task("tna", "merged")
        self._worktree("tna")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner(authed=False)
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tna", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "gh_not_authenticated")

    def test_push_failure_returns_502_and_stays_merged(self):
        self._insert_task("tpf", "merged")
        self._worktree("tpf")
        self.server.pr_git_runner = self._git_runner(push_fail=True)
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tpf", self.owner_key)
        self.assertEqual(status, 502, body)
        self.assertEqual(body["step"], "push")
        self.assertEqual(self._task_row("tpf")[0], "merged")

    def test_success_writes_back_pr_and_audit(self):
        self._insert_task("tok", "merged")
        self._worktree("tok")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tok", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["pr_url"],
                         "https://example.com/repo/pull/42")
        self.assertEqual(body["pr_number"], 42)
        self.assertEqual(body["branch"], "host/tok")
        row = self._task_row("tok")
        self.assertEqual(row[0], "merged")
        self.assertEqual(row[1], "https://example.com/repo/pull/42")
        self.assertEqual(row[2], 42)
        self.assertIn("pr_created", self._audit_actions())

    def test_pr_module_auto_pr_gate(self):
        with self.assertRaises(PrSkipped) as ctx:
            create_task_pr("x", 0)
        self.assertEqual(ctx.exception.reason, "auto_pr_disabled")


class ProjectTaskFieldTestCase(_ServerCase):
    def test_create_task_executor_validation(self):
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "x",
             "executor": "bogus"},
            self.owner_key)
        self.assertEqual(status, 400, body)
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "x",
             "executor": "docker"},
            self.owner_key)
        self.assertEqual(status, 201, body)
        self.assertEqual(body["executor"], "docker")


class WorkerTablePrecisionTestCase(unittest.TestCase):
    def test_workers_time_columns_are_double_precision(self):
        # 回归：2026-10-26 主机实测，workers.last_seen 用 REAL 时，
        # Postgres 的 float4 把 epoch 秒舍入到 128 秒一档，存活判断
        # 误差可达 ±64 秒、甚至出现负 age。新表必须是 DOUBLE PRECISION。
        import sqlite3 as _sqlite3
        tmp = tempfile.TemporaryDirectory()
        db_path = os.path.join(tmp.name, "w.db")
        from cloud.workers import ensure_workers_table
        ensure_workers_table(db_path)
        conn = _sqlite3.connect(db_path)
        try:
            rows = conn.execute("PRAGMA table_info(workers)").fetchall()
        finally:
            conn.close()
            tmp.cleanup()
        by_name = {r[1]: r[2] for r in rows}
        self.assertIn("DOUBLE", by_name["last_seen"])
        self.assertIn("DOUBLE", by_name["created_at"])


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