#!/usr/bin/env python3
"""Test suite for deporder.  Stdlib only:  python3 test_deporder.py

Covers: spec validation, unknown-dependency reporting, cycle detection (with
cycle path), deterministic topological order, parallel waves, weighted
critical path, greedy multi-worker schedule invariants, differential testing
against naive O(V*E) reference implementations on 200 random DAGs, the CLI
contract (exit codes 0/2/3/4), and a 20k-node performance budget.
"""

from __future__ import annotations

import json
import os
import random
import subprocess
import sys
import tempfile
import time
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import deporder as D  # noqa: E402

DEPORDER = os.path.join(HERE, "deporder.py")
EPS = 1e-9


# --------------------------------------------------------------- helpers


def tasks_from(deps, est=None):
    est = est or {}
    return {n: {"deps": list(deps[n]), "est": est.get(n, 1.0)} for n in deps}


def ref_waves(deps):
    """O(V*E) fixpoint reference for wave assignment."""
    level = {n: 0 for n in deps}
    changed = True
    while changed:
        changed = False
        for n in sorted(deps):
            want = 0 if not deps[n] else 1 + max(level[d] for d in deps[n])
            if want != level[n]:
                level[n] = want
                changed = True
    buckets = {}
    for n, lv in level.items():
        buckets.setdefault(lv, []).append(n)
    return [sorted(buckets[k]) for k in sorted(buckets)]


def ref_order(deps):
    """O(V^2) reference: repeatedly emit the smallest fully-ready task."""
    done = []
    left = set(deps)
    while left:
        cand = sorted(n for n in left if all(d in done for d in deps[n]))
        if not cand:
            raise AssertionError("reference hit a cycle")
        done.append(cand[0])
        left.discard(cand[0])
    return done


def ref_dist(deps, est):
    """Bellman-Ford style fixpoint for longest weighted path from any root."""
    dist = {n: float(est[n]) for n in deps}
    changed = True
    while changed:
        changed = False
        for n in sorted(deps):
            want = float(est[n]) + max((dist[d] for d in deps[n]), default=0.0)
            if want > dist[n]:
                dist[n] = want
                changed = True
    return dist


def rand_dag(rng, n, max_deps=3):
    names = ["n%03d" % i for i in range(n)]
    deps = {}
    for i, nm in enumerate(names):
        pool = names[:i]
        k = rng.randint(0, min(max_deps, len(pool)))
        deps[nm] = sorted(rng.sample(pool, k))
    return deps


def check_schedule(tc, deps, est, workers, sched, makespan):
    """Assert every scheduling invariant that must hold for any valid plan."""
    names = sorted(deps)
    tc.assertEqual(sorted(r["task"] for r in sched), names, "each task scheduled exactly once")
    by = {r["task"]: r for r in sched}
    for r in sched:
        tc.assertGreaterEqual(r["end"] - r["start"], float(est[r["task"]]) - EPS)
    for n in names:
        for d in deps[n]:
            tc.assertLessEqual(by[d]["end"], by[n]["start"] + EPS, "precedence %s<%s" % (d, n))
    per_worker = {}
    for r in sched:
        per_worker.setdefault(r["worker"], []).append(r)
    for w, rs in per_worker.items():
        rs.sort(key=lambda r: r["start"])
        for a, b in zip(rs, rs[1:]):
            tc.assertLessEqual(a["end"], b["start"] + EPS, "worker %d overlap" % w)
    events = []
    for r in sched:
        events.append((r["start"], 1))
        events.append((r["end"], -1))
    events.sort()
    live = 0
    for _, delta in events:
        live += delta
        tc.assertLessEqual(live, workers, "more than %d tasks in flight" % workers)
    tc.assertAlmostEqual(makespan, max(r["end"] for r in sched), places=9)
    cp_len, _ = D.critical_path(deps, {n: float(est[n]) for n in deps})
    tc.assertGreaterEqual(makespan, cp_len - EPS, "makespan below critical path")


def run_cli(args, spec=None):
    path = None
    tmpdir = None
    if spec is not None:
        tmpdir = tempfile.mkdtemp()
        path = os.path.join(tmpdir, "spec.json")
        with open(path, "w", encoding="utf-8") as fh:
            fh.write(spec if isinstance(spec, str) else json.dumps(spec))
    cmd = [sys.executable, DEPORDER] + list(args) + ([path] if path else [])
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if tmpdir:
        try:
            os.remove(path)
            os.rmdir(tmpdir)
        except OSError:
            pass
    return proc


# ------------------------------------------------------------------ tests


class SpecTests(unittest.TestCase):
    def test_defaults(self):
        tasks, workers = D.parse_spec({"tasks": {"a": {}, "b": {"deps": ["a"]}}})
        self.assertEqual(workers, 1)
        self.assertEqual(tasks["a"], {"deps": [], "est": 1.0})
        self.assertEqual(tasks["b"], {"deps": ["a"], "est": 1.0})

    def test_workers_override(self):
        _, workers = D.parse_spec({"tasks": {"a": {}}, "workers": 3})
        self.assertEqual(workers, 3)

    def test_duplicate_deps_collapse(self):
        tasks, _ = D.parse_spec({"tasks": {"a": {}, "b": {"deps": ["a", "a", "a"]}}})
        self.assertEqual(tasks["b"]["deps"], ["a"])

    def test_est_coerced_to_float(self):
        tasks, _ = D.parse_spec({"tasks": {"a": {"est": 2}}})
        self.assertIsInstance(tasks["a"]["est"], float)

    def test_rejects_malformed(self):
        bad = [
            None, [], "x", {},
            {"tasks": []},
            {"tasks": {}},
            {"tasks": {"a": []}},
            {"tasks": {"a": {"deps": "b"}}},
            {"tasks": {"a": {"deps": [1]}}},
            {"tasks": {"a": {"deps": [""]}}},
            {"tasks": {"a": {"est": 0}}},
            {"tasks": {"a": {"est": -1}}},
            {"tasks": {"a": {"est": True}}},
            {"tasks": {"a": {"est": "1"}}},
            {"tasks": {"a": {"est": float("inf")}}},
            {"tasks": {"a": {}}, "workers": 0},
            {"tasks": {"a": {}}, "workers": -2},
            {"tasks": {"a": {}}, "workers": 2.5},
            {"tasks": {"a": {}}, "workers": True},
        ]
        for spec in bad:
            with self.assertRaises(D.SpecError, msg="should reject %r" % (spec,)):
                D.parse_spec(spec)


class GraphTests(unittest.TestCase):
    def test_unknown_dep_reports_sorted_pairs(self):
        spec = {"tasks": {"b": {"deps": ["zz"]}, "a": {"deps": ["yy"]}}}
        tasks, _ = D.parse_spec(spec)
        with self.assertRaises(D.UnknownDepError) as ctx:
            D.build_graph(tasks)
        self.assertEqual(ctx.exception.pairs, [("a", "yy"), ("b", "zz")])
        self.assertIn("a -> yy", str(ctx.exception))

    def test_acyclic_has_no_cycle(self):
        deps = {"a": [], "b": ["a"], "c": ["a", "b"]}
        self.assertIsNone(D.find_cycle(deps))

    def test_self_loop(self):
        deps = {"a": ["a"]}
        self.assertEqual(D.find_cycle(deps), ["a"])
        with self.assertRaises(D.CycleError) as ctx:
            D.topological_order(deps)
        self.assertEqual(ctx.exception.cycle, ["a"])

    def test_two_cycle(self):
        deps = {"a": ["b"], "b": ["a"]}
        cycle = D.find_cycle(deps)
        self.assertEqual(len(cycle), 2)
        self.assertEqual(set(cycle), {"a", "b"})
        self.assertIn(cycle[0], deps[cycle[1]])

    def test_three_cycle_path_is_a_real_cycle(self):
        deps = {"a": ["c"], "b": ["a"], "c": ["b"]}
        cycle = D.find_cycle(deps)
        self.assertEqual(len(cycle), 3)
        for i, n in enumerate(cycle):
            self.assertIn(cycle[i - 1], deps[n], "edge %s -> %s missing" % (cycle[i - 1], n))
        with self.assertRaises(D.CycleError):
            D.waves(deps)

    def test_cycle_in_one_component_only(self):
        deps = {"a": [], "b": ["a"], "x": ["y"], "y": ["x"]}
        self.assertEqual(set(D.find_cycle(deps)), {"x", "y"})
        with self.assertRaises(D.CycleError):
            D.topological_order(deps)

    def test_cycle_error_message(self):
        deps = {"a": ["b"], "b": ["a"]}
        with self.assertRaises(D.CycleError) as ctx:
            D.critical_path(deps, {"a": 1.0, "b": 1.0})
        msg = str(ctx.exception)
        self.assertTrue(msg.startswith("dependency cycle: "), msg)
        head, tail = msg.split(": ", 1)
        parts = tail.split(" -> ")
        self.assertEqual(parts[0], parts[-1], "cycle path must return to its start")


class OrderTests(unittest.TestCase):
    def test_diamond_is_lexicographic(self):
        deps = {"d": ["b", "c"], "c": ["a"], "b": ["a"], "a": []}
        self.assertEqual(D.topological_order(deps), ["a", "b", "c", "d"])

    def test_chain(self):
        deps = {"c": ["b"], "b": ["a"], "a": []}
        self.assertEqual(D.topological_order(deps), ["a", "b", "c"])

    def test_disconnected_components_interleave(self):
        deps = {"a1": [], "b1": [], "a2": ["a1"], "b2": ["b1"]}
        self.assertEqual(D.topological_order(deps), ["a1", "a2", "b1", "b2"])
        self.assertEqual(D.topological_order(deps), ref_order(deps))

    def test_order_prefix_property_on_random_dags(self):
        rng = random.Random(20261021)
        for _ in range(200):
            deps = rand_dag(rng, rng.randint(1, 12))
            order = D.topological_order(deps)
            seen = set()
            for n in order:
                self.assertTrue(set(deps[n]) <= seen, "dep of %s emitted late" % n)
                seen.add(n)
            self.assertEqual(order, ref_order(deps))


class WaveTests(unittest.TestCase):
    def test_diamond(self):
        deps = {"d": ["b", "c"], "c": ["a"], "b": ["a"], "a": []}
        self.assertEqual(D.waves(deps), [["a"], ["b", "c"], ["d"]])

    def test_chain(self):
        deps = {"c": ["b"], "b": ["a"], "a": []}
        self.assertEqual(D.waves(deps), [["a"], ["b"], ["c"]])

    def test_parallel_roots_and_skew(self):
        deps = {"a": [], "b": [], "c": ["a"], "d": ["c"], "e": ["a", "b"]}
        self.assertEqual(D.waves(deps), [["a", "b"], ["c", "e"], ["d"]])

    def test_single_node(self):
        self.assertEqual(D.waves({"solo": []}), [["solo"]])

    def test_waves_match_reference_and_cover_every_task(self):
        rng = random.Random(7)
        for _ in range(200):
            deps = rand_dag(rng, rng.randint(1, 14))
            wv = D.waves(deps)
            self.assertEqual(wv, ref_waves(deps))
            self.assertEqual(sorted(x for w in wv for x in w), sorted(deps))
            for w in wv:
                self.assertEqual(w, sorted(w))


class CriticalPathTests(unittest.TestCase):
    def test_known_length(self):
        deps = {"a": [], "b": ["a"], "c": ["a"], "d": ["b", "c"]}
        est = {"a": 3.0, "b": 1.0, "c": 4.0, "d": 2.0}
        length, path = D.critical_path(deps, est)
        self.assertAlmostEqual(length, 3.0 + 4.0 + 2.0, places=9)
        self.assertEqual(path, ["a", "c", "d"])
        self.assertAlmostEqual(sum(est[n] for n in path), length, places=9)

    def test_path_is_precedence_chained(self):
        deps = {"a": [], "b": ["a"], "c": ["a", "b"], "d": ["b"]}
        est = {"a": 1.0, "b": 1.0, "c": 5.0, "d": 1.0}
        _, path = D.critical_path(deps, est)
        for prev, nxt in zip(path, path[1:]):
            self.assertIn(prev, deps[nxt])

    def test_tie_break_is_lexicographic(self):
        deps = {"a": [], "b": [], "z": ["a", "b"]}
        est = {"a": 2.0, "b": 2.0, "z": 1.0}
        length, path = D.critical_path(deps, est)
        self.assertAlmostEqual(length, 3.0, places=9)
        self.assertEqual(path, ["a", "z"])

    def test_single_node(self):
        self.assertEqual(D.critical_path({"a": []}, {"a": 2.5}), (2.5, ["a"]))

    def test_matches_reference(self):
        rng = random.Random(99)
        for _ in range(200):
            deps = rand_dag(rng, rng.randint(1, 14))
            est = {n: float(rng.choice([0.5, 1, 2, 3, 4])) for n in deps}
            length, path = D.critical_path(deps, est)
            self.assertAlmostEqual(length, max(ref_dist(deps, est).values()), places=9)
            self.assertAlmostEqual(sum(est[n] for n in path), length, places=9)


class ScheduleTests(unittest.TestCase):
    DEPS = {"a": [], "b": ["a"], "c": ["a"], "d": ["b", "c"], "e": ["d"]}
    EST = {"a": 3.0, "b": 1.0, "c": 4.0, "d": 2.0, "e": 1.0}

    def test_single_worker_equals_serial(self):
        sched = D.schedule(self.DEPS, self.EST, 1)
        check_schedule(self, self.DEPS, self.EST, 1, sched, max(r["end"] for r in sched))
        self.assertAlmostEqual(max(r["end"] for r in sched), sum(self.EST.values()), places=9)

    def test_unbounded_workers_equals_critical_path(self):
        w = len(self.DEPS)
        sched = D.schedule(self.DEPS, self.EST, w)
        check_schedule(self, self.DEPS, self.EST, w, sched, max(r["end"] for r in sched))
        cp, _ = D.critical_path(self.DEPS, self.EST)
        self.assertAlmostEqual(max(r["end"] for r in sched), cp, places=9)

    def test_two_workers_never_exceed_capacity(self):
        sched = D.schedule(self.DEPS, self.EST, 2)
        check_schedule(self, self.DEPS, self.EST, 2, sched, max(r["end"] for r in sched))
        self.assertAlmostEqual(max(r["end"] for r in sched), 10.0, places=9)

    def test_empty(self):
        self.assertEqual(D.schedule({}, {}, 3), [])

    def test_schedule_invariants_on_random_dags(self):
        rng = random.Random(4242)
        for _ in range(150):
            deps = rand_dag(rng, rng.randint(1, 12))
            est = {n: float(rng.choice([1, 2, 3, 5])) for n in deps}
            workers = rng.randint(1, 4)
            sched = D.schedule(deps, est, workers)
            check_schedule(self, deps, est, workers, sched, max(r["end"] for r in sched))


class PlanTests(unittest.TestCase):
    SPEC = {
        "tasks": {"a": {"est": 3}, "b": {"deps": ["a"], "est": 1},
                  "c": {"deps": ["a"], "est": 4}, "d": {"deps": ["b", "c"], "est": 2}},
        "workers": 2,
    }

    def test_plan_fields(self):
        p = D.plan(self.SPEC)
        self.assertEqual(p["order"], ["a", "b", "c", "d"])
        self.assertEqual(p["waves"], [["a"], ["b", "c"], ["d"]])
        self.assertEqual(p["critical_path"]["tasks"], ["a", "c", "d"])
        self.assertAlmostEqual(p["critical_path"]["length"], 9.0, places=9)
        self.assertAlmostEqual(p["serial_time"], 10.0, places=9)
        self.assertAlmostEqual(p["makespan"], 9.0, places=9)

    def test_key_insertion_order_does_not_change_output(self):
        base = D.plan(self.SPEC)
        shuffled = dict(reversed(list(self.SPEC["tasks"].items())))
        alt = D.plan({"tasks": shuffled, "workers": 2})
        self.assertEqual(json.dumps(base, sort_keys=True), json.dumps(alt, sort_keys=True))

    def test_plan_is_repeatable(self):
        a = json.dumps(D.plan(self.SPEC), sort_keys=True)
        b = json.dumps(D.plan(self.SPEC), sort_keys=True)
        self.assertEqual(a, b)


class CliTests(unittest.TestCase):
    def test_ok_exit_0(self):
        proc = run_cli(["plan"], PlanTests.SPEC)
        self.assertEqual(proc.returncode, 0, proc.stderr)
        out = json.loads(proc.stdout)
        self.assertEqual(out["order"], ["a", "b", "c", "d"])
        self.assertEqual(out["workers"], 2)

    def test_workers_flag_overrides_spec(self):
        spec = {"tasks": {"a": {}, "b": {}}}
        proc = run_cli(["plan", "--workers", "2"], spec)
        self.assertEqual(proc.returncode, 0, proc.stderr)
        self.assertEqual(json.loads(proc.stdout)["workers"], 2)

    def test_cycle_exit_2(self):
        spec = {"tasks": {"a": {"deps": ["b"]}, "b": {"deps": ["a"]}}}
        proc = run_cli(["plan"], spec)
        self.assertEqual(proc.returncode, 2, proc.stdout)
        self.assertIn("dependency cycle", proc.stderr)
        self.assertEqual(proc.stdout.strip(), "")

    def test_unknown_dep_exit_3(self):
        proc = run_cli(["plan"], {"tasks": {"a": {"deps": ["ghost"]}}})
        self.assertEqual(proc.returncode, 3)
        self.assertIn("unknown dependencies: a -> ghost", proc.stderr)

    def test_bad_spec_exit_4(self):
        proc = run_cli(["plan"], {"tasks": {"a": {"est": -1}}})
        self.assertEqual(proc.returncode, 4)
        proc = run_cli(["plan"], "{not json")
        self.assertEqual(proc.returncode, 4)
        self.assertIn("invalid JSON", proc.stderr)

    def test_missing_file_exit_4(self):
        proc = subprocess.run([sys.executable, DEPORDER, "plan", "/nonexistent/spec.json"],
                              capture_output=True, text=True)
        self.assertEqual(proc.returncode, 4)

    def test_bad_subcommand_exit_4(self):
        proc = run_cli(["frobnicate"], {"tasks": {"a": {}}})
        self.assertEqual(proc.returncode, 4)


class PerfTests(unittest.TestCase):
    def test_20k_node_graph_under_budget(self):
        n = 20000
        deps = {}
        est = {}
        for i in range(n):
            nm = "t%05d" % i
            deps[nm] = ["t%05d" % (i - 1)] if i else []
            est[nm] = 1.0
        rng = random.Random(1)
        names = sorted(deps)
        for _ in range(n):
            j = rng.randrange(1, n)
            i = rng.randrange(0, j)
            deps[names[j]].append(names[i])
        for nm in deps:
            deps[nm] = sorted(set(deps[nm]))

        t0 = time.perf_counter()
        order = D.topological_order(deps)
        wv = D.waves(deps)
        cp_len, cp_path = D.critical_path(deps, est)
        sched = D.schedule(deps, est, 8)
        elapsed = time.perf_counter() - t0

        self.assertEqual(len(order), n)
        self.assertEqual(sorted(x for w in wv for x in w), names)
        self.assertAlmostEqual(cp_len, float(len(cp_path)), places=9)
        self.assertEqual(len(sched), n)
        self.assertLess(elapsed, 10.0, "20k-node plan took %.2fs" % elapsed)
        sys.stderr.write("\n[perf] 20k nodes / %d edges: %.2fs (critical path %.0f)\n"
                         % (sum(len(v) for v in deps.values()), elapsed, cp_len))


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