import json
import os
import subprocess
import unittest

from test_cli import TOML, Cli
from test_claims import git
from test_merge import IDENT

TASKS = """\
# Tasks — demo

## Awaiting your decision

## Pending

- **t-one** [P1] (<1h): One.
  - Done: one works
  - Model: sonnet

- **t-two** [P2] (<1h): Two.
  - Done: two works

- **t-raw** [P2] (<1h): Raw, no Done line.

- **t-big** [P2] (1h): Big.
  - Done: big works

## Needs human

## Deferred
"""


class OrchTest(Cli):
    tasks_text = TASKS
    toml = TOML.replace('verify = ["make test"]\n', "")

    def setUp(self):
        super().setUp()
        git(self.root, "init", "-q", "-b", "master")
        (self.root / ".gitignore").write_text(".worktrees/\n.wf/\nout/\n")
        git(self.root, "add", "-A")
        git(self.root, "commit", "-qm", "init")

    def orch(self, *args, code=0):
        got, out, err = self.wf("orch", *args, env=IDENT)
        self.assertEqual(got, code, out + err)
        self.assertEqual(err, "")
        return out

    def run_wf(self, cwd, *args):
        got, out, err = self.wf(*args, project=False, cwd=cwd, env=IDENT)
        self.assertEqual(got, 0, out + err)
        return out

    def record(self, id):
        return json.loads((self.root / ".wf" / "orch" / f"{id}.json").read_text())

    def worker_done(self, id="t-one", lane="fast", merge=True):
        """What a worker does: wf start, code, wf done/finish (finish merges, done alone leaves the branch)."""
        wt = self.root / ".worktrees" / lane
        self.run_wf(self.root, "start", id, "--worktree", str(wt), "--branch", f"{lane}/{id}")
        (wt / "code.txt").write_text("x\n")
        if merge:
            self.run_wf(wt, "finish", id, "-m", "ok", "--commit", "impl", "code.txt", "--no-push")
        else:
            git(wt, "add", "code.txt")
            git(wt, "commit", "-qm", "impl")
            self.run_wf(wt, "done", id, "-m", "ok")
        return wt

    def log(self):
        return (self.root / "out" / "wf-orch.log").read_text()

    def test_pick_claims_and_prints_prompt(self):
        out = self.orch("pick", "fast")
        wt = self.root / ".worktrees" / "fast"
        self.assertEqual(out, "pick: t-one (lane fast, model sonnet, effort <1h) · claimed (in progress: worker)\n"
                              "agent: subagent_type wf-worker, model sonnet, no isolation; prompt:\n"
                              "Task: t-one  Lane: fast  Model: sonnet\n"
                              f"Main tree: {self.root}  Worktree: {wt}  Branch: fast/t-one\n"
                              "Final message: the 4 report lines only.\n")
        self.assertIn("- **t-one** [P1] (<1h) (in progress: worker): One.", self.tasks())
        self.assertEqual(self.record("t-one")["worktree"], str(wt))

    def test_second_pick_skips_in_progress_and_busy_worktree(self):
        self.orch("pick", "fast")
        out = self.orch("pick", "fast")
        self.assertIn("pick: t-two (lane fast, model opus", out)
        self.assertIn(f"Worktree: {self.root / '.worktrees' / 'fast-2'}  Branch: fast/t-two\n", out)

    def test_worktree_on_other_branch_not_reused(self):
        git(self.root, "worktree", "add", "-q", str(self.root / ".worktrees" / "fast"), "-b", "fast/t-old")
        self.assertIn(".worktrees/fast-2  Branch: fast/t-one", self.orch("pick", "fast"))

    def test_clean_detached_worktree_reused(self):
        git(self.root, "worktree", "add", "-q", "--detach", str(self.root / ".worktrees" / "fast"), "master")
        self.assertIn(".worktrees/fast  Branch: fast/t-one", self.orch("pick", "fast"))

    def test_worktree_with_live_session_not_reused(self):
        wt = self.root / ".worktrees" / "fast"
        git(self.root, "worktree", "add", "-q", "--detach", str(wt), "master")
        folder = self.root.parent / "no-claude" / "sessions"
        folder.mkdir(parents=True)
        (folder / "1.json").write_text(json.dumps({"pid": os.getpid(), "cwd": str(wt / "docs")}))
        self.assertIn(".worktrees/fast-2  Branch: fast/t-one", self.orch("pick", "fast"))

    def test_stop_file_spawns_nothing(self):
        (self.root / "out").mkdir()
        (self.root / "out" / "wf-batch.stop").write_text("")
        self.assertEqual(self.orch("pick", "fast"), "stop: out/wf-batch.stop exists: spawn nothing (let running workers finish)\n")
        self.assertNotIn("in progress", self.tasks())

    def test_none_pickable(self):
        self.orch("pick", "fast")
        self.orch("pick", "fast")
        self.assertTrue(self.orch("pick", "fast").startswith("none: lane fast has no runner-ready task"))

    def test_explicit_id_and_recovery_line(self):
        out = self.orch("pick", "slow", "--id", "t-big", "--recovery", "crashed")
        self.assertIn("Branch: slow/t-big\nRecovery: crashed\nFinal message", out)

    def test_post_done_logs_and_picks_next(self):
        self.orch("pick", "fast")
        self.worker_done()
        out = self.orch("post", "t-one", "fast", "--result", "done", "--commit", "abc1234", "--duration", "75")
        self.assertTrue(out.startswith("post: t-one done\n\npick: t-two"), out)
        self.assertRegex(self.log(), r"^\S+ fast sonnet t-one done abc1234 1m15s\n$")
        self.assertFalse((self.root / ".wf" / "orch" / "t-one.json").exists())

    def test_post_no_next_claims_nothing(self):
        self.orch("pick", "fast")
        self.worker_done()
        out = self.orch("post", "t-one", "fast", "--result", "done", "--no-next", "--duration", "5")
        self.assertTrue(out.startswith("post: t-one done"), out)
        self.assertNotIn("pick:", out)
        self.assertEqual([f.name for f in (self.root / ".wf" / "orch").glob("*.json")], [])

    def test_post_done_without_record_takes_model_from_history(self):
        self.orch("pick", "fast")
        self.worker_done()
        (self.root / ".wf" / "orch" / "t-one.json").unlink()
        self.orch("post", "t-one", "fast", "--result", "done", "--no-pick", "--duration", "5")
        self.assertRegex(self.log(), r"^\S+ fast sonnet t-one done ")

    def test_post_done_merges_unmerged_branch(self):
        self.orch("pick", "fast")
        wt = self.worker_done(merge=False)
        out = self.orch("post", "t-one", "fast", "--result", "done", "--no-pick", "--no-push")
        self.assertEqual(out, "post: t-one done\n  merged worktree HEAD: merged fast/t-one into master\n")
        self.assertEqual(subprocess.run(["git", "-C", str(self.root), "show", "master:code.txt"],
                                        capture_output=True, text=True).stdout, "x\n")
        self.assertEqual(subprocess.run(["git", "-C", str(wt), "branch", "--show-current"],
                                        capture_output=True, text=True).stdout, "")

    def test_post_done_without_archive_is_red(self):
        self.orch("pick", "fast")
        out = self.orch("post", "t-one", "fast", "--result", "done")
        self.assertTrue(out.startswith(
            "post: t-one post-check-red\n  no archive line for t-one\n  committed leftover TASKS.md tasks/archive.md\n"
            "alert: t-one post-check-red → tell the owner (Sessions: owner); lane fast keeps picking\n\n"
            "pick: t-two "), out)
        self.assertNotIn("stop lane", out)
        self.assertIn(" post-check-red ", self.log())
        self.assertIn("  Sessions: owner\n", self.wf("show", "t-one")[1])

    def backdate_pick(self, id="t-one", secs=125):
        f = self.root / ".wf" / "orch" / f"{id}.json"
        rec = json.loads(f.read_text())
        rec["at"] -= secs
        f.write_text(json.dumps(rec) + "\n")

    def test_post_zero_duration_falls_back_to_pick_time(self):
        self.orch("pick", "fast")
        self.backdate_pick()
        self.worker_done()
        self.orch("post", "t-one", "fast", "--result", "done", "--no-pick", "--duration", "0")
        self.assertRegex(self.log(), r" t-one done \S+ 2m0[5-9]s\n$")

    def test_post_red_keeps_pick_time_for_the_repost(self):
        self.orch("pick", "fast")
        self.backdate_pick()
        self.orch("post", "t-one", "fast", "--result", "done")
        self.assertTrue((self.root / ".wf" / "orch" / "t-one.json").exists())
        self.worker_done()
        self.orch("post", "t-one", "fast", "--result", "done", "--no-pick")
        self.assertRegex(self.log().splitlines()[-1], r" t-one done \S+ 2m0[5-9]s$")
        self.assertFalse((self.root / ".wf" / "orch" / "t-one.json").exists())

    def test_post_handback_parks_task_and_picks_next(self):
        self.orch("pick", "fast")
        out = self.orch("post", "t-one", "fast", "--result", "handback unclear spec")
        self.assertTrue(out.startswith(
            "post: t-one handback\n  committed leftover TASKS.md tasks/archive.md\n"
            "alert: t-one handback → tell the owner (Sessions: owner, status cleared); lane fast keeps picking\n\n"
            "pick: t-two (lane fast"), out)
        self.assertNotIn("stop lane", out)
        shown = self.wf("show", "t-one")[1]
        self.assertIn("  Sessions: owner\n", shown)
        self.assertIn("handback (unclear spec)", shown)
        self.assertNotIn("in progress", shown.splitlines()[0])
        self.assertEqual(subprocess.run(["git", "-C", str(self.root), "log", "-1", "--format=%s"],
                                        capture_output=True, text=True).stdout, "t-one handback (orchestrator)\n")
        self.assertIn("Sessions: owner", subprocess.run(["git", "-C", str(self.root), "show", "HEAD:TASKS.md"],
                                                        capture_output=True, text=True).stdout)

    def test_post_awaiting_blocks_task_and_picks_next(self):
        self.orch("pick", "fast")
        out = self.wf("add", "-s", "awaiting", "Which way?")[1]
        aid = next(w.strip("*[]:") for w in out.split() if w.strip("*[]:").startswith("a-"))
        out = self.orch("post", "t-one", "fast", "--result", f"awaiting {aid}", "--no-pick")
        self.assertEqual(out, f"post: t-one awaiting\n  committed leftover TASKS.md tasks/archive.md\n"
                              f"alert: t-one awaiting → tell the owner (blocked on {aid}); lane fast keeps picking\n")
        self.assertIn(f"blocked: [[{aid}]]", self.wf("show", "t-one")[1])
        self.assertIn("pick: t-two", self.orch("pick", "fast"))

    def test_post_needs_owner_waits_on_human_task_until_done(self):
        self.orch("pick", "fast")
        hid = "t-plug"
        self.assertEqual(self.wf("add", "-s", "human", "-p", "1", "-e", "<1h", "--id", hid, "Plug in the device. Owner-only step.")[0], 0)
        out = self.orch("post", "t-one", "fast", "--result", f"needs-owner {hid}", "--no-pick")
        self.assertEqual(out, f"post: t-one needs-owner\n  committed leftover TASKS.md tasks/archive.md\n"
                              f"alert: t-one needs-owner → tell the owner (after {hid}); lane fast keeps picking\n")
        shown = self.wf("show", "t-one")[1]
        self.assertIn(f"After: [[{hid}]]", shown)
        self.assertNotIn("in progress", shown.splitlines()[0])
        self.assertIn("pick: t-two", self.orch("pick", "fast"))       # A not picked while H open
        self.assertEqual(self.wf("done", hid, "-m", "plugged")[0], 0)
        self.assertIn("pick: t-one", self.orch("pick", "fast"))       # H done → A picked again

    def test_post_done_on_slice_job_counts_as_sliced(self):
        self.orch("pick", "slow", "--id", "t-big")
        self.run_wf(self.root, "add", "--parent", "t-big", "-e", "<1h", "--done", "x", "Slice one")
        out = self.orch("post", "t-big", "slow", "--result", "done", "--no-next", "--duration", "5")
        self.assertTrue(out.startswith("post: t-big sliced"), out)
        self.assertNotIn("post-check-red", self.log())
        self.assertIn(" t-big sliced ", self.log())

    def test_post_result_prefix_stripped(self):
        self.orch("pick", "fast")
        out = self.orch("post", "t-one", "fast", "--result", "result: handback unclear spec")
        self.assertTrue(out.startswith("post: t-one handback\n"), out)
        self.assertNotIn("stop lane", out)

    def test_post_result_prefix_done(self):
        self.orch("pick", "fast")
        self.worker_done()
        out = self.orch("post", "t-one", "fast", "--result", "result: done", "--no-pick")
        self.assertTrue(out.startswith("post: t-one done"), out)

    def test_post_model_raised_continues(self):
        self.orch("pick", "fast")
        self.wf("set", "t-one", "--model", "opus")
        self.wf("status", "t-one", "clear")
        out = self.orch("post", "t-one", "fast", "--result", "handback beyond sonnet")
        self.assertTrue(out.startswith("post: t-one model-raised\n"), out)
        self.assertIn("pick: t-one (lane fast, model opus", out)

    def test_post_no_report_suggests_recovery(self):
        self.orch("pick", "fast")
        self.assertIn('wf orch pick fast --id t-one --recovery "<why>"', self.orch("post", "t-one", "fast"))

    def test_post_agent_without_transcript_still_logs(self):
        self.orch("pick", "fast")
        out = self.orch("post", "t-one", "fast", "--result", "wip", "--agent", "abc")
        self.assertIn("  cost: not logged (no transcript for agent abc)\n", out)

    def test_refused_in_linked_worktree(self):
        git(self.root, "worktree", "add", "-q", "--detach", str(self.root / ".worktrees" / "x"), "master")
        got, _, err = self.wf("orch", "pick", "fast", project=False, cwd=self.root / ".worktrees" / "x")
        self.assertEqual((got, err), (1, "wf: orch runs in the main tree (the orchestrator's)\n"))


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