import json
import os
import sys
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from wflib import lanes as L
from wflib import tasks as T
from test_cli import Cli

ITEMS = """\
## Awaiting your decision

## Pending

- **t-a** [P1] (1h): A.

- **t-b** [P1] (1h): B.
  - After: [[t-s]]

- **t-s** [P1] (1h): S.
  Model: sonnet

- **t-s2** [P2] (1h): S2.
  Model: sonnet
  - After: [[t-x]]

- **t-c** [P2] (1h): C.
  - After: [[t-s2]]

## Needs human

## Deferred
"""
SOCK = "/run/user/1000/cc-socks/1.sock"


class LanesTest(unittest.TestCase):
    def setUp(self):
        self.doc = T.parse(ITEMS)

    def test_counts(self):
        self.assertEqual(L.counts(self.doc, set(), L.DEFAULT_LANES, "1h"),
                         {"fast": (0, 0, 0), "slow": (2, 3, 0)})

    def test_counts_held(self):
        self.assertEqual(L.counts(self.doc, set(), L.DEFAULT_LANES, "1h", {"t-a": "x"}),
                         {"fast": (0, 0, 0), "slow": (1, 4, 0)})

    def test_counts_slice_jobs(self):
        self.assertEqual(L.counts(T.parse(PICK), {"t-old1"}, L.DEFAULT_LANES, "1h"),
                         {"fast": (3, 2, 1), "slow": (1, 1, 0)})

    def test_lanes_block(self):
        counts = {"fast": (3, 1, 1), "slow": (0, 2, 0)}
        sessions = {"slow": {"socket": SOCK, "alive": True, "model": "opus"}}
        self.assertEqual(L.lanes_block(counts, sessions, "fast"), [
            "fast (you): 3 pickable (1 slice) · 1 waiting",
            f"slow: 0 pickable · 2 waiting · session uds:{SOCK} (alive, opus)",
        ])

    def test_lanes_block_no_session(self):
        self.assertEqual(L.lanes_block({"fast": (2, 0, 2), "slow": (0, 0, 0)}, {}, None), [
            "fast: 2 pickable (2 slices) · 0 waiting · no session → orchestrator or owner: start one (wf next --lane fast)",
            "slow: 0 pickable · 0 waiting · no session",
        ])

    def test_cross_waits_by_lane(self):
        doc = T.parse(PICK)
        self.assertEqual([(m.id, b.id) for m, b in L.cross_waits(doc, {"t-old1"}, "slow", L.DEFAULT_LANES, "1h")],
                         [("t-w1", "t-blocker")])
        self.assertEqual(L.cross_waits(doc, {"t-old1"}, "fast", L.DEFAULT_LANES, "1h"), [])

    def test_waiting_block(self):
        doc = T.parse(PICK)
        waits = L.cross_waits(doc, {"t-old1"}, "slow", L.DEFAULT_LANES, "1h")
        self.assertEqual(L.waiting_block(waits, {"fast": {"socket": SOCK, "alive": True}}, L.DEFAULT_LANES, "1h"),
                         [f'- t-w1 waits on t-blocker (fast lane) → message uds:{SOCK}: '
                          '"t-blocker blocks my t-w1, please take it"'])
        self.assertEqual(L.waiting_block(waits, {}, L.DEFAULT_LANES, "1h"),
                         ["- t-w1 waits on t-blocker (fast lane) → no fast session: tell the owner"])

    def test_notify_block_by_lane(self):
        doc = T.parse(PICK)
        freed = [doc.item("t-w1"), doc.item("t-w2")]
        self.assertEqual(L.notify_block(freed, {"slow": {"socket": SOCK, "alive": True}}, L.DEFAULT_LANES, "1h"), [
            "fast work now pickable: t-w2 (no session: tell the owner)",
            f"notify slow uds:{SOCK}: now pickable t-w1",
        ])

    def test_solo_done_block_any_session_name(self):
        sessions = {"fast": {"socket": SOCK, "alive": True, "pid": 1}, "all": {"socket": "/b", "alive": True, "pid": 2},
                    "slow": {"socket": "/c", "alive": False, "pid": 3}}
        self.assertEqual(L.solo_done_block(["t-s"], sessions, "2"),
                         [f"notify fast uds:{SOCK}: solo t-s done, run wf next"])


PICK = """\
## Awaiting your decision

## Pending

- **t-big** [P2] (5h): Big unsliced.

- **t-p0** [P0] (<1h): Small P0.

- **t-blocker** [P3] (<1h): Small, two wait on it.

- **t-w1** [P2] (1h): Waits.
  - After: [[t-blocker]]

- **t-w2** [P3] (<1h): Waits too.
  - After: [[t-blocker]]
  Model: sonnet

- **t-mid** [P1] (1h): Mid.
  Model: sonnet

- **t-done-parent** [P1] (10h): Sliced, slices archived.
  - Slices: [[t-old1]]

## Needs human

## Deferred
"""


class PickTest(unittest.TestCase):
    def setUp(self):
        self.doc = T.parse(PICK)

    def pick(self, lane=None, model=None, archived=frozenset({"t-old1"})):
        item, _ = L.pick(self.doc, set(archived), L.DEFAULT_LANES, "1h", lane, model)
        return item.id if item else None

    def test_lane_of(self):
        got = {i.id: L.lane_of(i, L.DEFAULT_LANES, "1h") for i in self.doc.section("pending").items}
        self.assertEqual(got, {"t-big": "fast", "t-p0": "fast", "t-blocker": "fast", "t-w1": "slow",
                               "t-w2": "fast", "t-mid": "slow", "t-done-parent": "fast"})

    def test_no_effort_is_fast(self):
        item = T.Item(id="t-x", prio=1)
        self.assertEqual(L.lane_of(item, L.DEFAULT_LANES, "1h"), "fast")

    def test_fast_unblock_order(self):
        # unblock: waited-on first (t-blocker: 2 waiters), then the slice job t-big (0 waiters) vs t-p0 by prio
        self.assertEqual([i.id for i in L.ranked(self.doc, {"t-old1"}, L.DEFAULT_LANES, "1h", "fast")],
                         ["t-blocker", "t-p0", "t-big"])

    def test_slow_priority_order_then_fallback(self):
        self.assertEqual(self.pick("slow"), "t-mid")
        self.assertEqual([i.id for i in L.ranked(self.doc, {"t-old1"}, L.DEFAULT_LANES, "1h", "slow")],
                         ["t-mid", "t-blocker", "t-p0", "t-big"])

    def test_fallback_when_lane_empty(self):
        doc = T.parse(PICK.replace("(1h): Mid.", "(<1h): Mid."))  # slow has only t-w1 (blocked)
        item, _ = L.pick(doc, {"t-old1"}, L.DEFAULT_LANES, "1h", "slow")
        self.assertEqual(item.id, "t-blocker")

    def test_all_lanes_priority_order(self):
        self.assertEqual(self.pick(None), "t-p0")

    def test_model_ceiling(self):
        # haiku session: no task has Model haiku; opus (default) and sonnet are above it
        self.assertIsNone(self.pick("fast", "haiku"))
        self.assertEqual(self.pick("slow", "sonnet"), "t-mid")
        self.assertEqual(self.pick("fast", "sonnet"), None)   # t-w2 is blocked; others are opus
        self.assertEqual(self.pick("fast", "opus"), "t-blocker")

    def test_slice_job(self):
        big = self.doc.item("t-big")
        self.assertTrue(L.is_slice_job(big, "1h"))
        self.assertFalse(L.is_slice_job(self.doc.item("t-p0"), "1h"))

    def test_sliced_parent_not_slice_job(self):
        parent = self.doc.item("t-done-parent")
        self.assertFalse(L.is_slice_job(parent, "1h"))
        self.assertTrue(L.sliced_out(parent, "1h"))
        # skipped lists items ranked before the pick: drop the others so the parent is first
        doc = T.parse(PICK[:PICK.index("- **t-big**")] + PICK[PICK.index("- **t-done-parent**"):])
        _, skipped = L.pick(doc, {"t-old1"}, L.DEFAULT_LANES, "1h", "fast")
        self.assertIn(("t-done-parent", "sliced: all slices done → wf done t-done-parent or add slices"),
                      [(i.id, why) for i, why in skipped])


PICK_DONE = PICK.replace("Small P0.\n", "Small P0.\n  Done: x\n").replace(
    "two wait on it.\n", "two wait on it.\n  Done: x\n").replace("Big unsliced.\n", "Big unsliced.\n  Done: x\n")


class LanesCliTest(Cli):
    tasks_text = "# Tasks — demo\n\n" + PICK_DONE

    def env(self, name="me"):
        sock = self.root / f"{name}.sock"
        sock.write_text("")
        return {"CLAUDE_CODE_MESSAGING_SOCKET": str(sock), "CLAUDE_PID": str(os.getpid()), "CLAUDE_CODE_SESSION_ID": name}

    def session(self, name):
        return json.loads((self.root / ".wf" / "sessions" / f"{name}.json").read_text())

    def test_next_lane_registers_lane(self):
        out = self.ok("next", "--lane", "slow", "--as", "opus", env=self.env())
        self.assertIn("===== Next task =====\n- **t-mid**", out)
        s = self.session("slow")
        self.assertEqual((s["lane"], s["model"], s["socket"]), ("slow", "opus", str(self.root / "me.sock")))
        self.assertIn("===== Lanes =====\nfast: 3 pickable (1 slice) · 2 waiting · no session → "
                      "orchestrator or owner: start one (wf next --lane fast)\n"
                      "slow (you): 1 pickable · 1 waiting\n\n"
                      "===== Waiting on other lanes =====\n"
                      "- t-w1 waits on t-blocker (fast lane) → no fast session: tell the owner\n", out)

    def test_next_without_lane_registers_all(self):
        out = self.ok("next", "--as", "opus", env=self.env())
        self.assertIn("===== Next task =====\n- **t-p0**", out)
        self.assertEqual(self.session("all")["lane"], "all")
        self.assertNotIn("Waiting on other lanes", out)

    def test_slice_job_output(self):
        (self.root / "TASKS.md").write_text("# Tasks — demo\n\n## Awaiting your decision\n\n## Pending\n\n"
                                            "- **t-big** [P2] (5h): Big unsliced.\n\n## Needs human\n\n## Deferred\n")
        out = self.ok("next", "--lane", "fast", "--as", "opus", "--brief")
        self.assertEqual(out, "- **t-big** [P2] (5h): Big unsliced.\n\n"
                              "===== Slice job (no code) =====\n"
                              "effort 5h > slice_above 1h: split it, don't implement.\n"
                              '  wf add "<title>. <goal>" -e <1h|1h> --parent t-big --model haiku|sonnet|opus'
                              "  (then wf body: Steps/Done/Ref)\n"
                              '  wf note t-big "sliced into …" · wf status t-big clear · not wf done'
                              " (the parent waits on its slices)\n")

    def test_unknown_lane(self):
        self.assertEqual(self.fails("next", "--lane", "nope", code=2), "wf: unknown lane 'nope' (fast, slow)\n")
        self.assertEqual(self.fails("list", "--lane", "nope", code=2), "wf: unknown lane 'nope' (fast, slow)\n")

    def test_nothing_pickable_text(self):
        self.assertEqual(self.fails("next", "--lane", "fast", "--as", "haiku", "--brief"),
                         "wf: nothing pickable for fast (haiku) in Pending\n")
        self.assertEqual(self.fails("next", "--as", "haiku", "--brief").splitlines()[-1],
                         "wf: nothing pickable for all lanes (haiku) in Pending")

    def test_list_runner_lane_pick_order(self):
        out = self.ok("list", "--runner", "--lane", "fast")
        self.assertEqual([l.split()[0] for l in out.splitlines()[:-1]], ["t-blocker", "t-p0", "t-big"])

    def test_list_lane_filter_and_column(self):
        out = self.ok("list", "--lane", "slow")
        self.assertEqual(out.splitlines()[:2], ["t-w1   P2  1h    -     opus    slow  Waits",
                                                "t-mid  P1  1h    -     sonnet  slow  Mid"])

    def test_lanes_command(self):
        out = self.ok("lanes", "--lane", "fast", "--as", "sonnet", env=self.env())
        self.assertEqual(out, "fast (you): 3 pickable (1 slice) · 2 waiting\n"
                              "slow: 0 pickable · 1 waiting · 1 not runner-ready (no Done) · no session\n")
        self.assertEqual((self.session("fast")["lane"], self.session("fast")["model"]), ("fast", "sonnet"))
        self.ok("lanes", "--unregister", env=self.env())
        self.assertFalse((self.root / ".wf" / "sessions" / "fast.json").exists())

    def test_done_notifies_by_lane(self):
        sock = self.root / "s.sock"
        sock.write_text("")
        d = self.root / ".wf" / "sessions"
        d.mkdir(parents=True)
        (d / "slow.json").write_text(json.dumps({"lane": "slow", "model": "opus", "socket": str(sock),
                                                 "pid": os.getpid()}))
        out = self.ok("done", "t-blocker", "-m", "ok", env={"CLAUDE_CODE_MESSAGING_SOCKET": ""})
        self.assertIn(f"notify slow uds:{sock}: now pickable t-w1\n", out)
        self.assertNotIn("t-w2", out)   # same lane as t-blocker: the doer sees it




WAIT_NONE = """\
## Awaiting your decision

## Pending

- **t-x** [P1] (1h): X.
  - After: [[t-nope]]

## Needs human

## Deferred
"""


class LanesWaitTest(Cli):
    tasks_text = WAIT_NONE
    env = {"WF_LANES_POLL": "0.1"}

    def test_pickable_exits_0_at_once(self):
        (self.root / "TASKS.md").write_text(DONE_ITEMS)
        out = self.ok("lanes", "--wait", "5", env=self.env)
        self.assertIn("slow: 2 pickable", out)

    def test_timeout_exits_1(self):
        self.assertEqual(self.fails("lanes", "--wait", "1", env=self.env), "")

    def test_added_mid_wait(self):
        import threading
        import time

        def add():
            time.sleep(0.6)
            (self.root / "TASKS.md").write_text(DONE_ITEMS)
        th = threading.Thread(target=add)
        th.start()
        out = self.ok("lanes", "--wait", "10", env=self.env)
        th.join()
        self.assertIn("slow: 2 pickable", out)

    def test_stop_file_exits_2_at_once(self):
        (self.root / "out").mkdir()
        (self.root / "out" / "wf-batch.stop").touch()
        code, out, err = self.wf("lanes", "--wait", "30", env=self.env)
        self.assertEqual((code, err), (2, ""))
        self.assertIn("stop requested: out/wf-batch.stop", out)

    def test_no_done_not_pickable(self):
        (self.root / "TASKS.md").write_text(NO_DONE)
        code, out, err = self.wf("lanes", "--wait", "1", env=self.env)
        self.assertEqual((code, err), (1, ""))
        self.assertEqual(out,
                         "fast: 0 pickable · 0 waiting · 2 not runner-ready (no Done) · no session\n"
                         "slow: 0 pickable · 0 waiting · no session\n")

    def test_lanes_plain_shows_not_ready(self):
        (self.root / "TASKS.md").write_text(NO_DONE)
        self.assertIn("fast: 0 pickable · 0 waiting · 2 not runner-ready (no Done)", self.ok("lanes"))


NO_DONE = """\
## Awaiting your decision

## Pending

- **t-a** [P1] (<1h): A.

- **t-b** [P1] (<1h): B.

## Needs human

## Deferred
"""

DONE_ITEMS = ITEMS.replace("Model: sonnet\n", "Model: sonnet\n  Done: ok.\n").replace("(1h): A.", "(1h): A.\n  Done: ok.") \
    .replace("(1h): B.", "(1h): B.\n  Done: ok.")


class LibRunnerTest(unittest.TestCase):
    def test_counts_runner_and_not_ready(self):
        doc = T.parse(NO_DONE.replace("A.", "A.\n  Done: ok."))
        a = (doc, set(), L.DEFAULT_LANES, "1h")
        self.assertEqual(L.counts(*a), {"fast": (2, 0, 0), "slow": (0, 0, 0)})
        self.assertEqual(L.counts(*a, runner=True), {"fast": (1, 0, 0), "slow": (0, 0, 0)})
        self.assertEqual(L.not_ready(*a), {"fast": 1, "slow": 0})

    def test_not_ready_skips_owner_bound_with_done(self):
        doc = T.parse(NO_DONE.replace("A.", "A.\n  Done: ok.\n  Sessions: owner"))
        a = (doc, set(), L.DEFAULT_LANES, "1h")
        self.assertEqual(L.not_ready(*a), {"fast": 1, "slow": 0})
        self.assertEqual(len(L.prep_targets(*a)), 1)


PREP = """\
## Pending

- **t-ok** [P1] (<1h): Has Done.
  Done: x works.

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

- **t-wait** [P0] (1h): Waits, no Done.
  - After: [[t-low]]

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

- **t-own** [P0] (<1h): Owner, no Done.
  Sessions: owner

- **t-prog** [P0] (<1h) (in progress: w): Running.

- **t-blk** [P0] (<1h) (blocked: a-q): Blocked.

- **t-par** [P0] (5h): Parent.

- **t-par-1** [P2] (<1h): Slice.
  Done: y.

## Needs human

- **t-h** [P0] (<1h): Human.
"""


class PrepTargetsTest(unittest.TestCase):
    def ids(self, **kw):
        return [i.id for i in L.prep_targets(T.parse(PREP), set(), L.DEFAULT_LANES, "1h", **kw)]

    def test_pickable_first_then_prio(self):
        self.assertEqual(self.ids(), ["t-hi", "t-low", "t-wait"])

    def test_lane_filter(self):
        self.assertEqual(self.ids(only={"fast"}), ["t-low"])
        self.assertEqual(self.ids(only={"slow"}), ["t-hi", "t-wait"])


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