"""U-013 tests: c_work_node_run's contract (schema validation, retry state
machine, quota-wait, execution_log_file lifecycle) with a fake `run_fn`, plus
real git worktree isolation + merge -- REQ-03's boundary case (two independent
nodes, separate worktrees, merge without conflict) needs no LLM call at all,
so it's tested for real here, not mocked.

Run: python3 -m unittest src.test_work_node_run -v
"""

import json
import os
import subprocess
import tempfile
import unittest

from src import schema, work_node_run

_IN_SCHEMA = {"type": "object", "required": ["goal"]}
_OUT_SCHEMA = {"type": "object", "required": ["value"]}


def _init_git_repo(path):
    subprocess.run(["git", "init"], cwd=path, check=True, capture_output=True)
    with open(os.path.join(path, "README.md"), "w") as f:
        f.write("seed\n")
    subprocess.run(["git", "add", "README.md"], cwd=path, check=True, capture_output=True)
    subprocess.run(
        ["git"] + work_node_run._GIT_IDENTITY + ["commit", "-m", "seed"],
        cwd=path, check=True, capture_output=True,
    )


class WorkNodeRunTestBase(unittest.TestCase):
    def setUp(self):
        self.conn = schema.connect(":memory:")
        schema.create_schema(self.conn)
        self.conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        self.conn.execute(
            "INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
            "VALUES ('task-1', 'proj-1', 'goal', 'goal', 'coding', 'coder')"
        )
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
            "VALUES ('run-1', 'task-1', 'coder', '{}', 'running')"
        )
        self.conn.execute(
            "INSERT INTO skill (skill_id, owner_org_node, file_path, current_git_commit_sha, source, status) "
            "VALUES ('skill-1', 'coder', 'skills/skill-1/', 'sha', 'new', 'active')"
        )
        self.conn.execute(
            "INSERT INTO work_node_execution (node_exec_id, run_id, node_id, subagent_type, skill_ref, "
            "input_schema_ref, output_schema_ref, input_payload, status) "
            "VALUES ('node-1', 'run-1', 'n1', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'pending')"
        )
        self.conn.commit()

    def tearDown(self):
        self.conn.close()


class InputValidationTest(WorkNodeRunTestBase):
    def test_invalid_input_rejected_before_execution_starts(self):
        with self.assertRaises(work_node_run.InvalidInputError):
            work_node_run.run_work_node(
                self.conn, "node-1", {}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=lambda cwd, model=None: {"value": 1}
            )
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "pending")
        count = self.conn.execute("SELECT COUNT(*) FROM execution_log_file").fetchone()[0]
        self.assertEqual(count, 0)


class HappyPathTest(WorkNodeRunTestBase):
    def test_valid_output_marks_node_done(self):
        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=lambda cwd, model=None: {"value": 42}
        )
        self.assertEqual(output, {"value": 42})
        row = self.conn.execute(
            "SELECT status, output_payload, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row[0], "done")
        self.assertEqual(json.loads(row[1]), {"value": 42})
        self.assertEqual(row[2], 0)

    def test_execution_log_file_created_active_then_closed(self):
        seen_status_during_run = []

        def run_fn(cwd, model=None):
            seen_status_during_run.append(
                self.conn.execute("SELECT status FROM execution_log_file").fetchone()[0]
            )
            return {"value": 1}

        work_node_run.run_work_node(self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn)
        self.assertEqual(seen_status_during_run, ["active"])
        final_status = self.conn.execute("SELECT status FROM execution_log_file").fetchone()[0]
        self.assertEqual(final_status, "closed")

    def test_log_file_closed_even_when_run_fn_raises(self):
        def boom(cwd, model=None):
            raise RuntimeError("node blew up")

        with self.assertRaises(RuntimeError):
            work_node_run.run_work_node(self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=boom)
        final_status = self.conn.execute("SELECT status FROM execution_log_file").fetchone()[0]
        self.assertEqual(final_status, "closed")


class RetryOnSchemaFailureTest(WorkNodeRunTestBase):
    def test_retries_then_succeeds(self):
        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] < 3:
                return {"wrong_key": True}
            return {"value": "ok"}

        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn, max_retries=2
        )
        self.assertEqual(output, {"value": "ok"})
        row = self.conn.execute(
            "SELECT status, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row, ("done", 2))

    def test_exceeding_max_retries_blocks_the_node(self):
        with self.assertRaises(work_node_run.SchemaValidationFailedError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
                run_fn=lambda cwd, model=None: {"wrong_key": True}, max_retries=2,
            )
        row = self.conn.execute(
            "SELECT status, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row, ("blocked", 3))


class UsageLimitWaitTest(WorkNodeRunTestBase):
    def setUp(self):
        super().setUp()
        self.sleeps = []
        self.clock = [0.0]

    def _sleep(self, seconds):
        self.sleeps.append(seconds)
        self.clock[0] += seconds

    def _now(self):
        return self.clock[0]

    def test_session_limit_then_recovers(self):
        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] == 1:
                work_node_run.raise_if_usage_limited("Error: You've hit your session limit")
            return {"value": "recovered"}

        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now,
        )
        self.assertEqual(output, {"value": "recovered"})
        self.assertEqual(self.sleeps, [5 * 60])

        row = self.conn.execute(
            "SELECT quota_limit_type FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row[0], "session")
        event = self.conn.execute(
            "SELECT event_type FROM execution_event WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(event[0], "quota_wait_started")

    def test_opus_quota_retries_immediately_on_downgraded_model_no_backoff(self):
        # REQ-14 boundary: "不走退避轮询，直接降级到非 Opus 模型重试" -- an actual
        # retry with a different model, not just detect-and-fail.
        calls = []

        def run_fn(cwd, model=None):
            calls.append(model)
            if len(calls) == 1:
                work_node_run.raise_if_usage_limited("You've hit your Opus limit")
            return {"value": "recovered on " + str(model)}

        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now,
        )
        self.assertEqual(calls, [None, "claude-sonnet-5"])
        self.assertEqual(output, {"value": "recovered on claude-sonnet-5"})
        self.assertEqual(self.sleeps, [])  # no backoff wait for opus
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "done")

    def test_opus_quota_custom_downgrade_model_is_used(self):
        calls = []

        def run_fn(cwd, model=None):
            calls.append(model)
            if len(calls) == 1:
                work_node_run.raise_if_usage_limited("You've hit your Opus limit")
            return {"value": "ok"}

        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now, opus_downgrade_model="claude-fable-5",
        )
        self.assertEqual(calls, [None, "claude-fable-5"])

    def test_opus_quota_persisting_after_downgrade_raises_timeout(self):
        # a non-opus model shouldn't hit an opus-specific quota again -- if it
        # somehow does, this doesn't loop forever chasing it
        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your Opus limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=self._sleep, now=self._now,
            )
        self.assertEqual(self.sleeps, [])

    def test_session_timeout_after_8_hours_blocks_node(self):
        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your session limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=self._sleep, now=self._now,
            )
        self.assertGreaterEqual(self.clock[0], 8 * 3600)
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "blocked")

    def test_weekly_limit_then_recovers(self):
        # previously untested quota_limit_type -- session/opus were covered,
        # weekly wasn't, even though the detection regex and timeout table
        # both already supported it
        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] == 1:
                work_node_run.raise_if_usage_limited("Error: You've hit your weekly limit")
            return {"value": "recovered"}

        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now,
        )
        quota_type = self.conn.execute(
            "SELECT quota_limit_type FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()[0]
        self.assertEqual(quota_type, "weekly")

    def test_weekly_timeout_after_9_days_blocks_node(self):
        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your weekly limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=self._sleep, now=self._now,
            )
        self.assertGreaterEqual(self.clock[0], 9 * 24 * 3600)
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "blocked")


class RealGitWorktreeIsolationTest(WorkNodeRunTestBase):
    """REQ-03 boundary: two independent nodes, each in its own git worktree,
    merged back with no conflict. Real git, no LLM call needed."""

    def setUp(self):
        super().setUp()
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        _init_git_repo(self.tmp.name)
        self.conn.execute(
            "INSERT INTO work_node_execution (node_exec_id, run_id, node_id, subagent_type, skill_ref, "
            "input_schema_ref, output_schema_ref, input_payload, status) "
            "VALUES ('node-2', 'run-1', 'n2', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'pending')"
        )
        self.conn.commit()

    def _write_file_run_fn(self, filename, content):
        def run_fn(cwd, model=None):
            with open(os.path.join(cwd, filename), "w") as f:
                f.write(content)
            subprocess.run(["git", "add", filename], cwd=cwd, check=True, capture_output=True)
            subprocess.run(
                ["git"] + work_node_run._GIT_IDENTITY + ["commit", "-m", "add " + filename],
                cwd=cwd, check=True, capture_output=True,
            )
            return {"value": filename}

        return run_fn

    def test_two_independent_nodes_isolated_and_merged_without_conflict(self):
        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=self._write_file_run_fn("module_a.txt", "module a content"),
            needs_worktree=True, repo_path=self.tmp.name,
        )
        work_node_run.run_work_node(
            self.conn, "node-2", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=self._write_file_run_fn("module_b.txt", "module b content"),
            needs_worktree=True, repo_path=self.tmp.name,
        )

        self.assertTrue(os.path.isfile(os.path.join(self.tmp.name, "module_a.txt")))
        self.assertTrue(os.path.isfile(os.path.join(self.tmp.name, "module_b.txt")))

        row1 = self.conn.execute(
            "SELECT git_worktree_ref FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()[0]
        row2 = self.conn.execute(
            "SELECT git_worktree_ref FROM work_node_execution WHERE node_exec_id = 'node-2'"
        ).fetchone()[0]
        self.assertNotEqual(row1, row2)
        self.assertIn("node-1", row1)
        self.assertIn("node-2", row2)

        log = subprocess.run(
            ["git", "log", "--oneline"], cwd=self.tmp.name, check=True, capture_output=True, text=True
        )
        self.assertIn("module_a.txt", log.stdout)
        self.assertIn("module_b.txt", log.stdout)


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