import io
import shlex
import subprocess
import sys
import tempfile
import threading
import unittest
from contextlib import redirect_stderr, redirect_stdout
from pathlib import Path

HERE = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(HERE))
sys.path.insert(0, str(HERE / "tests"))

import wf_res  # noqa: E402
from test_res_io import Box, hist_rec, seed_history  # noqa: E402

OUT = "out/wf-batch-2026-10-01-1400.md"


class BatchBase(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.box = Box(Path(self._tmp.name))
        bin_ = self.box.tmp / "bin"
        bin_.mkdir()
        self.claude = bin_ / "claude"
        self.claude.write_text("#!/bin/sh\n")
        self.claude.chmod(0o755)
        self.box.caller["PATH"] = f"{bin_}:/usr/bin"

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

    def batch(self, *argv):
        out, err = io.StringIO(), io.StringIO()
        with redirect_stdout(out), redirect_stderr(err):
            code = wf_res.batch_main(list(argv), self.box.env)
        return code, out.getvalue(), err.getvalue()

    def unit_argv(self):
        return self.box.fake.ran("systemd-run")[0]


class Start(BatchBase):
    def test_starts_claude_print_unit(self):
        code, out, err = self.batch("2", "--lanes", "fast,slow")
        self.assertEqual((code, err), (0, ""))
        log = self.box.env.state / "logs" / "r-1.log"
        self.assertEqual(out, f"r-1 started; log {log}; ETA ~20:00 (estimate: not killed when over)\n"
                              f"summary: {OUT} · status: wf batch --status\n")
        argv = self.unit_argv()
        i = argv.index(str(self.claude))
        self.assertEqual(argv[i - 1], "sh")
        self.assertEqual(argv[i + 1:-1], ["-p", "--model", "opus", "--permission-mode", "auto",
                                          "--permission-prompts", "none"])
        self.assertIn("--setenv=CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS=21600000", argv)
        self.assertIn(f"MemoryMax={4 * 1024 ** 3}", argv)
        e = self.box.ledger().get("r-1")
        self.assertEqual((e.title, e.est_min, e.project), ("wf-batch", 360, "proj"))

    def test_prompt_from_template(self):
        self.batch("2", "--lanes", "fast,slow")
        prompt = self.unit_argv()[-1]
        self.assertIn("for up to 2 tasks (lanes: fast, slow)", prompt)
        self.assertIn(f"append one line per task to {OUT}", prompt)
        self.assertIn("/projects/public/workflow/docs/orchestrator.md", prompt.replace(str(HERE), "/projects/public/workflow"))
        self.assertIn("foreground", prompt)
        self.assertNotIn("{", prompt)

    def test_dry_run_writes_no_summary(self):
        self.batch("2", "--dry-run")
        self.assertFalse((self.box.env.root() / OUT).exists())

    def test_existing_summary_not_overwritten(self):
        f = self.box.env.root() / OUT
        f.parent.mkdir(exist_ok=True)
        f.write_text("live\n")
        self.batch("2")
        self.assertTrue(f.read_text().startswith("live\n# wf-batch"))

    def test_size_lane(self):
        code, out, err = self.batch("2", "--lanes", "fast", "--dry-run")
        self.assertEqual((code, err), (0, ""))
        self.assertIn("lanes: fast", out)

    def test_default_lanes(self):
        self.batch("3")
        self.assertIn("for up to 3 tasks (lanes: every lane with ready tasks, see wf lanes)", self.unit_argv()[-1])

    def test_for_sets_ceiling_and_model(self):
        self.batch("1", "--for", "2h", "--model", "sonnet", "--mem", "2G")
        argv = self.unit_argv()
        self.assertIn("--setenv=CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS=7200000", argv)
        self.assertEqual(argv[argv.index("--model") + 1], "sonnet")
        self.assertIn(f"MemoryMax={2 * 1024 ** 3}", argv)

    def test_no_claude(self):
        self.box.caller["PATH"] = "/nonexistent"
        code, out, err = self.batch("2")
        self.assertEqual((code, out, err), (1, "", "wf: claude not found on PATH\n"))
        self.assertEqual(self.box.fake.ran("systemd-run"), [])

    def test_bad_count(self):
        code, _, err = self.batch("0")
        self.assertEqual(code, 2)
        self.assertIn("N must be ≥ 1", err)

    def test_summary_header_written(self):
        code, out, err = self.batch("2", "--lanes", "fast")
        self.assertEqual(code, 0)
        text = (self.box.env.root() / OUT).read_text()
        self.assertEqual(text, "# wf-batch proj 2026-10-01 14:00 (N=2, lanes: fast)\n")

    def test_dry_run(self):
        code, out, err = self.batch("2", "--dry-run")
        self.assertEqual((code, err), (0, ""))
        self.assertIn("CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS=21600000 wf res run --mem 4G --for 6h --title wf-batch -- ", out)
        self.assertIn(f"{self.claude} -p --model opus", out)
        self.assertEqual(self.box.fake.ran("systemd-run"), [])

    def test_mem_scales_with_n(self):
        for args, mem in ((("4",), "6G"), (("1",), "4G"), (("9",), "8G"), (("4", "--mem", "3G"), "3G")):
            code, out, err = self.batch(*args, "--dry-run")
            self.assertIn(f"wf res run --mem {mem} --for", out, args)

    def test_defaults_from_history(self):
        # peaks 2,3,4 → p95 4 ×1.15 = 4.6 GB; durations 20,30,40 → p90 40 ×1.5 = 60 min; ceiling stays 6h
        seed_history(self.box, [hist_rec(f"r-{i}", title="wf-batch", peak=p, minutes=m)
                                for i, (p, m) in enumerate(((2.0, 20.0), (3.0, 30.0), (4.0, 40.0)))])
        code, out, err = self.batch("4", "--dry-run")
        self.assertEqual((code, err), (0, ""))
        self.assertIn("CEILING_MS=21600000 wf res run --mem 4.6G --for 60m --title wf-batch -- ", out)
        out = self.batch("4", "--mem", "3G", "--for", "2h", "--dry-run")[1]
        self.assertIn("CEILING_MS=7200000 wf res run --mem 3G --for 2h --title wf-batch -- ", out)

    def test_busy_exit_3(self):
        self.box.available_gb = 1
        code, out, _ = self.batch("2")
        self.assertEqual(code, 3)
        self.assertIn("busy", out)

    def test_force_past_unused_claim(self):
        self.box.wf("run", "--mem", "10G", "--for", "40m", "--title", "big", "--", "x")
        code, out, _ = self.batch("2")
        self.assertEqual(code, 3)
        code, out, err = self.batch("2", "--force")
        self.assertEqual(code, 0, (out, err))
        self.assertIn("r-2 started", out)


class Status(BatchBase):
    def test_none(self):
        code, out, err = self.batch("--status")
        self.assertEqual((code, out, err), (0, "no batch in this project\n", ""))

    def test_running_and_summary(self):
        self.batch("2")
        md = self.box.tmp / "proj" / OUT
        with md.open("a") as f:
            f.write("t-a sonnet done abc123\n")
        (md.parent / "wf-batch-2026-09-30-0100.md").write_text("old\n")
        code, out, _ = self.batch("--status")
        self.assertEqual(code, 0)
        lines = out.splitlines()
        self.assertTrue(lines[0].startswith('r-1 proj "wf-batch" running'), lines[0])
        self.assertIn(f"log {self.box.env.state / 'logs' / 'r-1.log'}", lines[1])
        self.assertEqual(lines[2:], [f"{OUT}:", "# wf-batch proj 2026-10-01 14:00 (N=2, lanes: every lane with ready tasks, see wf lanes)",
                                    "t-a sonnet done abc123"])

    def test_orch_log_progress(self):
        self.batch("2")
        out_dir = self.box.tmp / "proj" / "out"
        (out_dir / "wf-orch.log").write_text(
            "2026-10-01 13:00 fast sonnet t-old done aaa\n2026-10-01 14:05 fast sonnet t-new done bbb\n")
        _, out, _ = self.batch("--status")
        self.assertIn("finished since batch start", out)
        self.assertIn("t-new done bbb", out)
        self.assertNotIn("t-old", out)

    def test_stop_creates_file_and_status_reports(self):
        code, out, err = self.batch("--stop")
        self.assertEqual((code, err), (0, ""))
        f = self.box.tmp / "proj" / "out" / "wf-batch.stop"
        self.assertTrue(f.exists())
        self.assertIn("out/wf-batch.stop", out)
        code, out, _ = self.batch("--status")
        self.assertIn("stop requested: out/wf-batch.stop", out)

    def test_stale_stop_cleared_on_start(self):
        f = self.box.tmp / "proj" / "out" / "wf-batch.stop"
        f.parent.mkdir(exist_ok=True)
        f.touch()
        code, out, err = self.batch("2")
        self.assertEqual(code, 0)
        self.assertIn("stale out/wf-batch.stop cleared", err)

    def test_stop_consumed_on_exit(self):
        f = self.box.tmp / "proj" / "out" / "wf-batch.stop"
        f.parent.mkdir(exist_ok=True)
        orig = wf_res.cmd_run

        def fake(env, cfg, args):
            f.touch()
            return 0
        wf_res.cmd_run = fake
        try:
            code, _, _ = self.batch("2")
        finally:
            wf_res.cmd_run = orig
        self.assertEqual(code, 0)
        self.assertFalse(f.exists())

    def test_prompt_has_stop_check(self):
        code, out, _ = self.batch("2", "--dry-run")
        self.assertIn("out/wf-batch.stop", out)
        self.assertIn("stopped: stop file", out)

    def test_prompt_cloud_lane(self):
        code, out, _ = self.batch("2", "--dry-run")
        self.assertIn("wf orch pick cloud", out)
        self.assertIn("never run `wf cloud pull`", out)
        self.assertIn("sidecar pulls every 10 min", out)
        self.assertIn(f"-- {self.claude} -p", out)          # not a cloud project: plain claude -p job

    def test_prompt_stale_branch_once(self):
        code, out, _ = self.batch("2", "--dry-run")
        self.assertIn("pre-existing branch", out)
        self.assertIn("ONCE", out)

    def test_prompt_checks_stop_before_every_spawn(self):
        code, out, _ = self.batch("2", "--dry-run")
        self.assertIn("Before EVERY spawn", out)
        self.assertIn(f"test -e {self.box.tmp / 'proj' / 'out' / 'wf-batch.stop'}", out)
        self.assertNotIn("Each round: first", out)
        code, out, _ = self.batch("--stop")
        self.assertIn("before its next spawn", out)

    def test_n_required_without_status(self):
        code, _, err = self.batch()
        self.assertEqual(code, 2)


PREP_TASKS = """\
## Awaiting your decision

## Pending

- **t-ready** [P1] (<1h): Ready.
  Done: works.

- **t-low** [P3] (<1h): Low, no Done.

- **t-hi** [P1] (1h): High, no Done.

- **t-mid** [P2] (1h): Mid, no Done.

## Needs human

## Deferred
"""


class Prep(BatchBase):
    def setUp(self):
        super().setUp()
        root = self.box.env.root()
        (root / "workflow.toml").write_text('format = 1\ntasks = "TASKS.md"\narchive = "archive.md"\n')
        (root / "TASKS.md").write_text(PREP_TASKS)
        (root / "archive.md").write_text("# Archive\n")

    def test_prompt_lists_targets(self):
        code, out, err = self.batch("2", "--prep")
        self.assertEqual((code, err), (0, ""))
        prompt = self.unit_argv()[-1]
        self.assertIn("Tasks: t-hi t-mid.", prompt)
        self.assertIn("never implement", prompt)
        self.assertIn("wf set <id> --done", prompt)
        self.assertIn("wf add -s awaiting", prompt)
        self.assertIn(f"to {OUT}", prompt)
        for ph in ("never invert recorded original behaviour", "wf add --parent <id> slices",
                   "concrete files/runs", "[opus: <why>]", "[Model: <m>]"):
            self.assertIn(ph, prompt)
        self.assertNotIn("{", prompt)
        self.assertEqual((self.box.env.root() / OUT).read_text(),
                         "# wf-batch proj 2026-10-01 14:00 (prep K=2: t-hi t-mid)\n")
        self.assertEqual(self.box.ledger().get("r-1").title, "wf-batch")

    def test_lanes_filter(self):
        self.batch("5", "--prep", "--lanes", "fast")
        self.assertIn("Tasks: t-low.", self.unit_argv()[-1])

    def test_prep_alias_all(self):
        out, err = io.StringIO(), io.StringIO()
        with redirect_stdout(out), redirect_stderr(err):
            code = wf_res.prep_main(["all"], self.box.env)
        self.assertEqual((code, err.getvalue()), (0, ""))
        self.assertIn("Tasks: t-hi t-mid", self.unit_argv()[-1])

    def test_nothing_to_prep_starts_nothing(self):
        code, out, err = self.batch("3", "--prep", "--lanes", "nolane")
        self.assertEqual((code, out, err), (0, "prep: no pending task without Done (lanes: nolane); nothing started\n", ""))
        self.assertEqual(self.box.fake.ran("systemd-run"), [])
        self.assertFalse((self.box.env.root() / OUT).exists())

    def test_prep_mem_default_small(self):
        for args, mem in ((("3", "--prep"), "1G"), (("3", "--prep", "--mem", "2G"), "2G")):
            code, out, err = self.batch(*args, "--dry-run")
            self.assertIn(f"wf res run --mem {mem} --for", out, args)

    def test_dry_run(self):
        code, out, err = self.batch("1", "--prep", "--dry-run")
        self.assertEqual((code, err), (0, ""))
        self.assertIn("--title wf-batch -- ", out)
        self.assertIn("t-hi", out)
        self.assertEqual(self.box.fake.ran("systemd-run"), [])

    def test_no_project(self):
        (self.box.env.root() / "workflow.toml").unlink()
        code, out, err = self.batch("1", "--prep")
        self.assertEqual(code, 1)
        self.assertTrue(err.startswith("wf: "), err)


class Dispatch(unittest.TestCase):
    def test_wf_forwards_batch(self):
        r = subprocess.run([sys.executable, str(HERE / "wf.py"), "batch", "-h"], capture_output=True, text=True)
        self.assertEqual(r.returncode, 0)
        self.assertIn("usage: wf batch", r.stdout)

    def test_listed_in_wf_help(self):
        r = subprocess.run([sys.executable, str(HERE / "wf.py"), "-h"], capture_output=True, text=True)
        self.assertIn("batch", r.stdout)


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


FAKE_WF = """import sys, pathlib
d = pathlib.Path(sys.argv[0]).parent
with open(d / "calls.log", "a") as f:
    f.write(" ".join(sys.argv[1:]) + "\\n")
if sys.argv[1:3] == ["cloud", "pull"] and (d / "clear").exists():
    for rec in (d / ".wf" / "cloud").glob("*.json"):
        rec.unlink()
if sys.argv[1:3] == ["cloud", "pull"] and not (d / "pulled").exists():
    (d / "pulled").touch()
    print("t-a: done (session_1, $0.40 usage): merged abc1234")
    print("archive session_1 failed: x; archive by hand (wf cloud archive --ended)")
    print("report: commit abc1234")
    print("t-b: running (session_2, 12 min)")
    print("t-c: handback (session_3, $0.10 usage): cloud handback: stuck")
    print("t-d: pulled by another wf cloud pull, skipped")
    sys.exit(1)
"""


class Sidecar(unittest.TestCase):
    """wf batch on a cloud = true project: the job runs the orchestrator under the pull sidecar."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.root = Path(self._tmp.name)
        (self.root / "out").mkdir()
        self.summary = self.root / OUT
        self.summary.write_text("# wf-batch\n")
        self.fake = self.root / "fakewf.py"
        self.fake.write_text(FAKE_WF)

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

    def run_sidecar(self, child_s, every=0.2, rc=7, **kw):
        child = [sys.executable, "-c", f"import time, sys; time.sleep({child_s}); sys.exit({rc})"]
        out = io.StringIO()
        with redirect_stdout(out):
            code = wf_res.sidecar(child, self.root, self.summary, every, [sys.executable, str(self.fake)], **kw)
        return code, out.getvalue()

    def records(self, *ids):
        d = self.root / ".wf" / "cloud"
        d.mkdir(parents=True)
        for i in ids:
            (d / f"{i}.json").write_text("{}")

    def tail_line(self):
        return [ln[6:] for ln in self.summary.read_text().splitlines()[1:] if "tail" in ln]

    def test_tail_ends_on_empty_records(self):
        self.records("t-b")
        (self.root / "clear").touch()
        code, out = self.run_sidecar(0.05, every=0.1)
        self.assertEqual(code, 7)
        self.assertEqual(len([c for c in self.calls() if c.startswith("cloud pull")]), 1)
        self.assertIn("orch post t-a cloud --result done --no-pick --commit abc1234", self.calls())
        self.assertEqual(self.tail_line(), ["cloud tail ended: no cloud records left; 1 pulls; left: none (sidecar)"])

    def test_tail_ends_on_stop_file(self):
        self.records("t-b", "t-e")
        threading.Timer(0.5, (self.root / wf_res.STOP_FILE).touch).start()
        code, _ = self.run_sidecar(0.05, every=0.2)
        pulls = [c for c in self.calls() if c.startswith("cloud pull")]
        self.assertGreaterEqual(len(pulls), 1)            # pulled in the tail until the stop file
        self.assertEqual(self.tail_line(), [f"cloud tail ended: stop file; {len(pulls)} pulls; left: t-b t-e (sidecar)"])
        self.assertEqual(code, 7)

    def test_tail_ends_on_cap(self):
        self.records("t-b")
        shrunk = []
        code, out = self.run_sidecar(0.05, every=0.1, cap=0.35, shrink=lambda: shrunk.append(1) or "shrunk")
        self.assertEqual((code, shrunk), (7, [1]))       # reservation shrunk once, entering the tail
        self.assertIn("shrunk", out)
        pulls = [c for c in self.calls() if c.startswith("cloud pull")]
        self.assertGreaterEqual(len(pulls), 2)
        self.assertTrue(all(c.startswith("cloud pull") or c.startswith("orch post") for c in self.calls()))  # no picks
        self.assertEqual(self.tail_line(),
                         [f"cloud tail ended: {0.35 / 3600:g}h cap; {len(pulls)} pulls; left: t-b (sidecar)"])

    def test_no_tail_without_records(self):
        shrunk = []
        code, _ = self.run_sidecar(0.05, every=0.1, shrink=lambda: shrunk.append(1) or "")
        self.assertEqual((code, shrunk, self.calls(), self.tail_line()), (7, [], [], []))

    def test_shrink_reservation(self):
        b = BatchBase("run")
        b.setUp()
        try:
            b.batch("4", "--mem", "6G")
            msg = wf_res.shrink_reservation(b.box.env, "r-1")
            e = b.box.ledger().get("r-1")
            self.assertEqual((e.mem_gb, e.cpus, e.state), (0.2, 1, "running"))
            self.assertIn("r-1 reservation 6", msg)
            self.assertIn("not running", wf_res.shrink_reservation(b.box.env, "r-9"))
        finally:
            b.tearDown()

    def calls(self):
        f = self.root / "calls.log"
        return f.read_text().splitlines() if f.exists() else []

    def test_pulls_and_posts_while_child_lives(self):
        code, out = self.run_sidecar(1.0)
        self.assertEqual(code, 7)                          # the orchestrator's exit code
        calls = self.calls()
        pulls = [c for c in calls if c.startswith("cloud pull")]
        self.assertGreaterEqual(len(pulls), 2)            # repeats every interval while the child lives
        self.assertEqual(pulls[0], f"cloud pull --all --project {self.root}")
        self.assertEqual([c for c in calls if c.startswith("orch post")],
                         ["orch post t-a cloud --result done --no-pick --commit abc1234",
                          "orch post t-c cloud --result handback --no-pick"])
        lines = self.summary.read_text().splitlines()[1:]
        self.assertEqual([ln[6:] for ln in lines], ["cloud t-a done abc1234 (sidecar pull)",
                                                    "cloud t-c handback - (sidecar pull)"])
        self.assertIn("t-a: done", out)                   # pull output lands in the job log

    def test_stops_with_child(self):
        code, _ = self.run_sidecar(0.1, every=5, rc=0)
        self.assertEqual((code, self.calls()), (0, []))

    def test_stop_file_ends_pulls(self):
        (self.root / wf_res.STOP_FILE).touch()
        code, out = self.run_sidecar(0.7)
        self.assertEqual((code, self.calls()), (7, []))
        self.assertIn("sidecar: stop file, no more pulls", out)

    def test_cloud_project_job_wraps_claude(self):
        b = BatchBase("run")
        b.setUp()
        try:
            root = b.box.env.root()
            (root / "workflow.toml").write_text('format = 1\ntasks = "TASKS.md"\narchive = "a.md"\ncloud = true\n')
            code, out, _ = b.batch("2", "--dry-run")
            self.assertEqual(code, 0)
            argv = shlex.split(out.split(" -- ", 1)[1])
            self.assertEqual(argv[:9], [sys.executable, str(HERE / "wf_res.py"), "batch-sidecar", "--root", str(root),
                                        "--summary", OUT, "--every", "600"])
            self.assertEqual(argv[9:12], ["--", str(b.claude), "-p"])
        finally:
            b.tearDown()

    def test_main_entry(self):
        r = subprocess.run([sys.executable, str(HERE / "wf_res.py"), "batch-sidecar", "--root", str(self.root),
                            "--summary", OUT, "--every", "5", "--wf", f"{sys.executable} {self.fake}", "--",
                            sys.executable, "-c", "import sys; sys.exit(3)"], capture_output=True, text=True)
        self.assertEqual((r.returncode, r.stderr), (3, ""))


LOG_ROWS = "".join(f"2026-10-01T1{i}:00:00 slow opus t-{i} done abc {m}m00s\n" for i, m in enumerate((20, 25, 31)))


class Fit(BatchBase):
    def log(self, text):
        (self.box.env.root() / "out").mkdir(exist_ok=True)
        (self.box.env.root() / "out" / "wf-orch.log").write_text(text)

    def test_left_shrinks_k_and_sets_for(self):
        self.log(LOG_ROWS)                                     # p90 of 20,25,31 = 31 min
        code, out, err = self.batch("4", "--left", "1h40m", "--dry-run")
        self.assertEqual((code, err), (0, ""))
        self.assertIn("fit: 3 of 4 tasks (left 1h40m, task p90 31m from 3 runs)\n", out)
        self.assertIn("CEILING_MS=6000000 wf res run --mem 4.5G --for 1h40m --title wf-batch -- ", out)
        self.assertIn("for up to 3 tasks", out)

    def test_left_default_30m(self):
        code, out, err = self.batch("4", "--left", "45m", "--dry-run")
        self.assertIn("fit: 1 of 4 tasks (left 45m, task p90 30m default (0 runs))\n", out)
        self.assertIn("--for 45m", out)

    def test_left_too_short_starts_nothing(self):
        self.log(LOG_ROWS)
        code, out, err = self.batch("4", "--left", "30m")
        self.assertEqual((code, err), (0, ""))
        self.assertEqual(out, "fit: 0 of 4 tasks (left 30m, task p90 31m from 3 runs); nothing started\n")
        self.assertEqual(self.box.fake.ran("systemd-run"), [])

    def test_prompt_has_deadline_check(self):
        code, out, _ = self.batch("2", "--for", "2h", "--dry-run")
        self.assertIn("wf batch --time-left 2026-10-01T16:00+00:00", out)
        self.assertIn("stopped: deadline", out)

    def test_time_left(self):
        self.log(LOG_ROWS)
        code, out, err = self.batch("--time-left", "2026-10-01T15:00+00:00")
        self.assertEqual((code, out, err), (0, "time left 1h, task p90 31m (3 runs): spawn\n", ""))
        code, out, err = self.batch("--time-left", "2026-10-01T14:30+00:00")
        self.assertEqual((code, out), (0, "time left 30m, task p90 31m (3 runs): stop\n"))
        code, out, err = self.batch("--time-left", "2026-10-01T13:00+00:00")
        self.assertEqual((code, out), (0, "time left 0m, task p90 31m (3 runs): stop\n"))

    def test_time_left_bad(self):
        code, out, err = self.batch("--time-left", "soon")
        self.assertEqual(code, 1)
        self.assertIn("wf: --time-left", err)