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

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from wflib import tasks as T
from wflib import lanes as L
from test_cli import Cli
from test_check import Base, CLEAN
from wflib import check as K
from wflib import config as C

LANES = """\
# Tasks — demo

## Awaiting your decision

## Pending

- **t-one** [P1] (1h): One.
  - Steps: a
  Model: sonnet
  - After: [[t-done]]

- **t-two** [P2] (1h): Two.
  - Model: haiku

- **t-three** [P2] (<1h): Three.

## Needs human

## Deferred
"""


class ModelLineTest(unittest.TestCase):
    def setUp(self):
        self.doc = T.parse(LANES)

    def test_parsed_from_body_default_opus(self):
        self.assertEqual([i.model for i in self.doc.section("pending").items], ["sonnet", "haiku", "opus"])

    def test_extra_words_ignored(self):
        item = T.parse_block("- **t-x** [P1] (1h): X.\n  - Model: sonnet ok\n")
        self.assertEqual(item.model, "sonnet")

    def test_set_replaces_in_place(self):
        T.set_fields(self.doc, "t-one", model="opus")
        self.assertEqual(self.doc.item("t-one").body, ["  - Steps: a", "  Model: opus", "  - After: [[t-done]]"])

    def test_set_adds_before_after_and_ref(self):
        doc = T.parse(CLEAN)
        T.set_fields(doc, "t-one", model="haiku")
        self.assertEqual(doc.item("t-one").body,
                         ["  Model: haiku", "  - After: [[t-done]]", "  Ref: DESIGN.md#terrain, docs/specs/terrain.md"])

    def test_set_empty_removes(self):
        T.set_fields(self.doc, "t-two", model="")
        self.assertEqual(self.doc.item("t-two").body, [])
        self.assertEqual(self.doc.item("t-two").model, "opus")

    def test_set_bad_value(self):
        with self.assertRaisesRegex(T.TaskError, r"model 'gpt' \(want haiku, sonnet, opus\)"):
            T.set_fields(self.doc, "t-one", model="gpt")

    def test_note_goes_before_model_line(self):
        T.add_note(self.doc, "t-two", "why")
        self.assertEqual(self.doc.item("t-two").body, ["  - why", "  - Model: haiku"])


def doc(items):
    return T.parse("## Awaiting your decision\n\n## Pending\n\n" + items + "\n## Needs human\n\n## Deferred\n")


class PickOrderTest(unittest.TestCase):
    def pick(self, items, model="opus", archived=()):
        item, _ = L.pick(doc(items), set(archived), L.DEFAULT_LANES, "1h", None, model)
        return item.id if item else None

    def test_priority_inherited_from_waiting_task(self):
        self.assertEqual(self.pick("- **t-c** [P1] (1h): C.\n\n- **t-b** [P3] (1h): B.\n\n"
                                   "- **t-a** [P0] (1h): A.\n  - After: [[t-b]]\n"), "t-b")

    def test_inherited_transitively(self):
        self.assertEqual(self.pick("- **t-c** [P1] (1h): C.\n\n- **t-x** [P3] (1h): X.\n\n"
                                   "- **t-b** [P3] (1h): B.\n  - After: [[t-x]]\n\n"
                                   "- **t-a** [P0] (1h): A.\n  - After: [[t-b]]\n"), "t-x")

    def test_unblocking_only_low_work_does_not_jump(self):
        self.assertEqual(self.pick("- **t-c** [P1] (1h): C.\n\n- **t-b** [P3] (1h): B.\n\n"
                                   "- **t-a** [P3] (1h): A.\n  - After: [[t-b]]\n"), "t-c")

    def test_other_lane_waiting_first(self):
        self.assertEqual(self.pick("- **t-f** [P2] (1h): F.\n\n- **t-k** [P2] (1h): K.\n\n"
                                   "- **t-l** [P2] (1h): L.\n  - After: [[t-k]]\n\n- **t-g** [P2] (1h): G.\n\n"
                                   "- **t-h** [P2] (<1h): H.\n  - After: [[t-g]]\n"), "t-g")

    def test_more_waiting_first(self):
        self.assertEqual(self.pick("- **t-f** [P2] (1h): F.\n  - After: [[t-z]]\n\n- **t-g** [P2] (1h): G.\n\n"
                                   "- **t-i** [P2] (1h): I.\n  - After: [[t-g]]\n\n"
                                   "- **t-j** [P2] (1h): J.\n  - After: [[t-i]]\n\n"
                                   "- **t-z** [P2] (1h): Z.\n", archived=()), "t-g")

    def test_parent_waits_on_its_slices(self):
        self.assertEqual(self.pick("- **t-q** [P1] (1h): Q.\n\n- **t-p** [P0] (5h): P.\n  - Slices: [[t-p-1]]\n\n"
                                   "- **t-p-1** [P2] (1h): P1.\n"), "t-p-1")

    def test_lane_filter(self):
        items = ("- **t-a** [P0] (1h): A.\n\n- **t-b** [P1] (1h): B.\n  Model: sonnet\n\n"
                 "- **t-c** [P2] (1h): C.\n  Model: haiku\n")
        self.assertEqual([self.pick(items, m) for m in ("opus", "sonnet", "haiku", None)],
                         ["t-a", "t-b", "t-c", "t-a"])
        self.assertIsNone(self.pick("- **t-a** [P0] (1h): A.\n", "haiku"))

    def test_skipped_are_ranked_before_pick(self):
        d = doc("- **t-a** [P0] (1h) (blocked: [[a-k]]): A.\n\n- **t-s** [P0] (1h): S.\n  Model: sonnet\n\n"
                "- **t-b** [P1] (1h): B.\n\n- **t-c** [P2] (1h): C.\n  - After: [[t-b]]\n")
        item, skipped = L.pick(d, set(), L.DEFAULT_LANES, "1h", "slow", "opus")
        self.assertEqual((item.id, [(i.id, why) for i, why in skipped]), ("t-s", [("t-a", "blocked: a-k")]))
        item, skipped = L.pick(d, set(), L.DEFAULT_LANES, "1h", "slow", "haiku")
        self.assertEqual((item, skipped), (None, []))


class ModelCheckTest(Base):
    def test_values(self):
        self.pending("- **t-a** [P1] (1h): A.\n  Model: gpt\n\n"
                     "- **t-b** [P1] (1h): B.\n  - Model: sonnet ok\n\n"
                     "- **t-c** [P1] (1h): C.\n  Model: haiku\n  Model: opus\n\n"
                     "- **t-d** [P1] (1h): D.\n")
        errors, warnings = K.check(C.load(self.root))
        self.assertEqual([(p.id, p.message) for p in errors],
                         [("t-a", "Model 'gpt' (want haiku, sonnet, opus)"), ("t-c", "two Model lines")])
        self.assertEqual([(p.id, p.message) for p in warnings],
                         [("t-b", "Model line 'sonnet ok': write 'Model: sonnet' (wf set --model)")])


class ModelCliTest(Cli):
    tasks_text = LANES

    def test_list_column_and_filter(self):
        self.assertEqual(self.ok("list"),
                         "t-one    P1  1h    -     sonnet  slow  One\n"
                         "t-two    P2  1h    -     haiku   slow  Two\n"
                         "t-three  P2  <1h   -     opus    fast  Three\n"
                         "pending 3 · human 0 · awaiting 0 · deferred 0\n")
        self.assertEqual(self.ok("list", "--model", "opus"),
                         "t-three  P2  <1h   -     opus    fast  Three\n"
                         "pending 3 · human 0 · awaiting 0 · deferred 0\n")

    def test_add_and_set(self):
        self.assertEqual(self.ok("add", "Four. Goal.", "-p", "3", "-e", "1h", "--model", "haiku", "--ref", "DESIGN.md"),
                         "- **t-four** [P3] (1h): Four. Goal.\n")
        self.assertEqual(self.item("t-four"), "- **t-four** [P3] (1h): Four. Goal.\n  Model: haiku\n  Ref: DESIGN.md\n")
        self.ok("set", "t-four", "--model", "sonnet")
        self.assertEqual(self.item("t-four"), "- **t-four** [P3] (1h): Four. Goal.\n  Model: sonnet\n  Ref: DESIGN.md\n")
        self.ok("set", "t-four", "--model", "")
        self.assertEqual(self.item("t-four"), "- **t-four** [P3] (1h): Four. Goal.\n  Ref: DESIGN.md\n")

    def test_next_model_ceiling(self):
        # --as M takes tasks with Model ≤ M; no --lane = all lanes, priority order
        self.assertEqual(self.ok("next", "--as", "haiku", "--brief"), "- **t-two** [P2] (1h): Two.\n  - Model: haiku\n")
        self.assertEqual(self.ok("next", "--as", "opus", "--brief").splitlines()[0], "- **t-one** [P1] (1h): One.")
        self.assertEqual(self.ok("next", "--as", "sonnet", "--brief").splitlines()[0], "- **t-one** [P1] (1h): One.")
        self.assertEqual(self.ok("next", "--lane", "fast", "--as", "opus", "--brief"), "- **t-three** [P2] (<1h): Three.\n")
        self.ok("set", "t-two", "--model", "opus")
        self.assertEqual(self.fails("next", "--as", "haiku", "--brief").splitlines()[-1],
                         "wf: nothing pickable for all lanes (haiku) in Pending")

    def test_next_without_as_is_haiku_with_hint(self):
        out = self.ok("next", "--brief")
        self.assertEqual(out, "no --as: treated as haiku; pass --as haiku|sonnet|opus\n"
                              "- **t-two** [P2] (1h): Two.\n  - Model: haiku\n")

    def session_env(self, sock):
        sock.write_text("")
        return {"CLAUDE_CODE_MESSAGING_SOCKET": str(sock), "CLAUDE_PID": str(os.getpid()),
                "CLAUDE_CODE_SESSION_ID": "sess-1"}

    def dead_pid(self):
        p = subprocess.Popen(["true"])
        p.wait()
        return p.pid

    def test_next_registers_session(self):
        env = self.session_env(self.root / "me.sock")
        self.ok("next", "--as", "opus", "--brief", env=env)
        reg = json.loads((self.root / ".wf" / "sessions" / "all.json").read_text())
        self.assertEqual({k: reg[k] for k in ("lane", "model", "socket", "pid", "session")},
                         {"lane": "all", "model": "opus", "socket": str(self.root / "me.sock"), "pid": os.getpid(), "session": "sess-1"})
        self.assertEqual((self.root / ".wf" / ".gitignore").read_text(), "*\n")

    def test_no_register_without_as_or_env(self):
        none = {"CLAUDE_CODE_MESSAGING_SOCKET": "", "CLAUDE_PID": "", "CLAUDE_CODE_SESSION_ID": ""}
        self.ok("next", "--as", "opus", env=none)
        self.ok("next", env=self.session_env(self.root / "me.sock"))
        self.assertFalse((self.root / ".wf").exists())

    def register(self, lane, model, sock, pid):
        d = self.root / ".wf" / "sessions"
        d.mkdir(parents=True, exist_ok=True)
        (d / f"{lane}.json").write_text(json.dumps({"lane": lane, "model": model, "socket": str(sock), "pid": pid,
                                                     "session": "s", "at": "2026-10-04T10:00"}))

    def test_next_shows_lanes_and_waiting(self):
        self.ok("set", "t-three", "--after", "t-one")
        self.ok("set", "t-two", "--model", "opus")
        sock = self.root / "son.sock"
        sock.write_text("")
        self.register("slow", "sonnet", sock, os.getpid())
        self.register("fast", "haiku", self.root / "gone.sock", self.dead_pid())
        none = {"CLAUDE_CODE_MESSAGING_SOCKET": ""}
        code, out, err = self.wf("next", "--lane", "fast", "--as", "opus", env=none)   # fast: no fallback lane
        self.assertEqual((code, err), (1, "wf: nothing pickable for fast (opus) in Pending\n"))
        self.assertIn("===== Lanes =====\n"
                      "fast (you): 0 pickable · 1 waiting\n"
                      f"slow: 2 pickable · 0 waiting · session uds:{sock} (alive, sonnet)\n\n"
                      "===== Waiting on other lanes =====\n"
                      f'- t-three waits on t-one (slow lane) → message uds:{sock}: "t-one blocks my t-three, please take it"\n', out)
        self.assertEqual(self.ok("lanes", env=none),
                         "fast: 0 pickable · 1 waiting · no session\n"
                         f"slow: 0 pickable · 0 waiting · 2 not runner-ready (no Done) · session uds:{sock} (alive, sonnet)\n")

    def test_next_single_lane_prints_no_lanes_block(self):
        (self.root / "workflow.toml").write_text(self.toml + '[lanes.one]\nefforts = ["<1h", "1h"]\n')
        self.assertNotIn("Lanes", self.ok("next", "--as", "opus", env={"CLAUDE_CODE_MESSAGING_SOCKET": ""}))

    def test_done_notifies_other_lanes(self):
        self.ok("add", "Four.", "-p", "3", "-e", "<1h", "--model", "haiku", "--after", "t-three")
        self.ok("add", "Five.", "-p", "3", "-e", "1h", "--after", "t-three")
        self.ok("add", "Six.", "-p", "3", "-e", "1h", "--model", "sonnet", "--after", "t-three")
        sock = self.root / "h.sock"
        sock.write_text("")
        self.register("slow", "opus", sock, os.getpid())
        out = self.ok("done", "t-three", "-m", "ok")
        self.assertIn(f"notify slow uds:{sock}: now pickable t-five, t-six\n", out)
        self.assertNotIn("t-four", out)   # same lane as t-three (fast)

    def test_bad_model_is_usage_error(self):
        self.fails("add", "Four.", "-p", "3", "-e", "1h", "--model", "gpt", code=2)
        self.fails("set", "t-one", "--model", "gpt", code=2)


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