"""标记等待中父任务心跳单测：sweep 误杀回归。

根因：标记等待在 _run_attempt 返回后才跑，那时 _running 已被清掉，
run_loop 的心跳线程刷不到等候中的父任务；sweep 默认阈值 300 秒，
而标记等待标称可到 3600 秒，父任务等超 300 秒就被当死任务回收。
修法：轮询 / 等待循环每次迭代自己刷一次心跳（间隔至多 interval，
钳制后 <=60 秒，远小于 300 秒阈值），不依赖心跳线程。

不跑真实模型、不联网：用 init_db 建的临时 SQLite 库 + 最小任务行，
直接调 _deliver_marker_tools，不走 run_one 全量流程。
"""

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

try:
    from cloud.server import init_db, sweep_stale_tasks
    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.server import init_db, sweep_stale_tasks
    from cloud.worker import Worker

FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}


class MarkerWaitHeartbeatTestCase(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)
        # 父任务 PM：running，但心跳已陈旧到离 sweep 阈值只差 1 秒。
        # 不修时，等 2 秒后心跳就过期（301 > 300），sweep 会误杀；
        # 修了时，等待循环每次迭代刷心跳，sweep 应放过它。
        stale = time.time() - 299
        self._insert("PM", status="running", last_heartbeat=stale,
                     updated_at=stale)
        self._insert("E1", status="queued", parent="PM")

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

    def _insert(self, task_id, status="queued", parent=None,
                last_heartbeat=None, updated_at=None):
        now = time.time()
        conn = sqlite3.connect(self.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(?, 'p1', 'do work', ?, NULL, '[]', 0,"
                " 'tester', ?, ?, ?)",
                (task_id, status, now,
                 updated_at if updated_at is not None else now,
                 last_heartbeat))
            if parent is not None:
                conn.execute(
                    "UPDATE tasks SET parent_task_id = ? WHERE id = ?",
                    (parent, task_id))
            conn.commit()
        finally:
            conn.close()

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

    def _heartbeat(self, task_id):
        conn = sqlite3.connect(self.db_path)
        try:
            row = conn.execute(
                "SELECT status, last_heartbeat FROM tasks WHERE id = ?",
                (task_id,)).fetchone()
            return row[0], row[1]
        finally:
            conn.close()

    def test_poll_wait_refreshes_heartbeat_and_survives_sweep(self):
        worker = self._worker()
        worker._deliver_marker_tools(worker.get_task("PM"), {
            "status": "completed", "task_creates": [],
            "task_polls": [{"task_id": "E1", "interval_seconds": 1,
                            "timeout_seconds": 2, "chars": 2000}],
            "wait_threads": [], "task_cancels": [], "task_retries": []})
        status, heartbeat = self._heartbeat("PM")
        self.assertEqual(status, "running")
        # 等待循环刷过心跳：心跳应比陈旧的 -299 秒新得多。
        self.assertGreater(heartbeat, time.time() - 60)
        swept = sweep_stale_tasks(self.db_path, 300)
        self.assertNotIn("PM", swept)
        status, _ = self._heartbeat("PM")
        self.assertEqual(status, "running")

    def test_wait_threads_refreshes_heartbeat_and_survives_sweep(self):
        worker = self._worker()
        worker._deliver_marker_tools(worker.get_task("PM"), {
            "status": "completed", "task_creates": [],
            "task_polls": [],
            "wait_threads": [{"task_ids": ["E1"],
                              "interval_seconds": 1,
                              "timeout_seconds": 2, "chars": 2000}],
            "task_cancels": [], "task_retries": []})
        status, heartbeat = self._heartbeat("PM")
        self.assertEqual(status, "running")
        self.assertGreater(heartbeat, time.time() - 60)
        swept = sweep_stale_tasks(self.db_path, 300)
        self.assertNotIn("PM", swept)
        status, _ = self._heartbeat("PM")
        self.assertEqual(status, "running")


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