"""PR needs_local_worker 持久化单测（残留批次 A 任务 2）。

- 非 local worker 扫到 queued 的 pr 任务后，行状态持久化为
  needs_local_worker（重读仍在），且返回条目状态一致；
- 标记过的行仍可被 local worker 认领（claim 后 running，可
  正常 finish），不会被搁置；
- find_open_pr_job 仍把标记过的行算未完成（/pr 不重排）；
- 标记是 CAS：已 running 的行不被覆盖；重复扫不重复记。

测试用的 key、主机名、任务 id 全是假值；不联网、不碰生产。
"""

import os
import sys
import tempfile
import unittest
from pathlib import Path

try:
    from cloud.db import connect as _db_connect
    from cloud.pr_exec import (
        claim_pr_job, enqueue_pr_job, find_open_pr_job,
        mark_pr_job_needs_local_worker)
    from cloud.sandbox import DockerExecutor
    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.db import connect as _db_connect
    from cloud.pr_exec import (
        claim_pr_job, enqueue_pr_job, find_open_pr_job,
        mark_pr_job_needs_local_worker)
    from cloud.sandbox import DockerExecutor
    from cloud.worker import Worker


class PrNeedsLocalTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")

    def _enqueue_t1(self, host="host-a"):
        conn = _db_connect(self.db_path)
        try:
            enqueue_pr_job(conn, "t1", "p1", host)
        finally:
            conn.close()

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

    def _row(self, task_id):
        conn = _db_connect(self.db_path)
        try:
            return conn.execute(
                "SELECT id, status, error FROM pr_jobs"
                " WHERE task_id = ?", (task_id,)).fetchone()
        finally:
            conn.close()

    def test_docker_peek_persists_mark(self):
        worker = Worker(self.db_path, "/nonexistent/host.py",
                        env_key="TEST_KEY",
                        executor=DockerExecutor("/fake/host.py"))
        conn = _db_connect(self.db_path)
        try:
            enqueue_pr_job(conn, "t1", "p1", worker.host_name)
        finally:
            conn.close()
        jobs = worker._run_pr_jobs()
        self.assertEqual(len(jobs), 1)
        self.assertEqual(jobs[0]["status"], "needs_local_worker")
        self.assertEqual(jobs[0]["task_id"], "t1")
        # 重读：行状态已持久化，不只是返回值里有。
        row = self._row("t1")
        self.assertEqual(row[1], "needs_local_worker")
        self.assertIn("local worker", row[2])
        # 再扫一次：已标记的行不再是 queued，不重复返回、不重复记。
        self.assertEqual(worker._run_pr_jobs(), [])

    def test_marked_row_still_claimable_by_local(self):
        self._enqueue_t1()
        conn = _db_connect(self.db_path)
        try:
            job_id = conn.execute(
                "SELECT id FROM pr_jobs WHERE task_id = 't1'"
            ).fetchone()[0]
            self.assertTrue(mark_pr_job_needs_local_worker(
                conn, job_id, "needs local worker"))
            # 标记后仍算未完成：/pr 去重不受影响。
            self.assertIsNotNone(find_open_pr_job(conn, "t1"))
            # local worker 认领标记过的行。
            claimed = claim_pr_job(conn, "host-a")
            self.assertIsNotNone(claimed)
            self.assertEqual(claimed["task_id"], "t1")
            row = conn.execute(
                "SELECT status FROM pr_jobs WHERE task_id = 't1'"
            ).fetchone()
            self.assertEqual(row[0], "running")
        finally:
            conn.close()

    def test_mark_is_cas_on_queued_only(self):
        self._enqueue_t1()
        conn = _db_connect(self.db_path)
        try:
            claimed = claim_pr_job(conn, "host-a")
            self.assertIsNotNone(claimed)
            # 已 running：标记必须失败、不覆盖状态。
            self.assertFalse(mark_pr_job_needs_local_worker(
                conn, claimed["id"], "late mark"))
            row = conn.execute(
                "SELECT status FROM pr_jobs WHERE task_id = 't1'"
            ).fetchone()
            self.assertEqual(row[0], "running")
        finally:
            conn.close()