#!/usr/bin/env python3
"""deporder -- deterministic build-order planner for task dependency DAGs.

Stdlib only.  Spec (JSON):

    {
      "tasks": {
        "compile": {"deps": [],            "est": 3},
        "link":    {"deps": ["compile"],   "est": 1}
      },
      "workers": 2
    }

Semantics
---------
* Edge set E = {(d, n) : d in tasks[n].deps} -- "d runs before n".
* `est` is a positive finite duration estimate; default 1.0.
* Every ordering decision is deterministic: ties are broken by
  lexicographically smallest task name (heap order), so two runs over the
  same spec -- and runs over specs whose keys are inserted in a different
  order -- produce byte-identical output.

API
---
parse_spec(obj)                  -> (tasks, workers)
build_graph(tasks)               -> (deps, rdeps)          raises UnknownDepError
find_cycle(deps)                 -> cycle path or None
topological_order(deps)          -> list[str]              raises CycleError
waves(deps)                      -> list[list[str]]        raises CycleError
critical_path(deps, est)         -> (length, path)         raises CycleError
schedule(deps, est, workers)     -> list[dict]             raises CycleError
plan(obj)                        -> dict                   full plan
main(argv)                       -> int                    CLI entry point

CLI
---
    python3 deporder.py plan spec.json [--workers N]

Exit codes: 0 ok, 2 dependency cycle, 3 unknown dependency, 4 malformed spec.
"""

from __future__ import annotations

import heapq
import json
import sys

INF = float("inf")
__all__ = [
    "SpecError",
    "UnknownDepError",
    "CycleError",
    "parse_spec",
    "build_graph",
    "find_cycle",
    "topological_order",
    "waves",
    "critical_path",
    "schedule",
    "plan",
    "main",
]


class SpecError(ValueError):
    """Malformed spec (exit code 4)."""


class UnknownDepError(SpecError):
    """A dependency names a task that does not exist (exit code 3)."""

    def __init__(self, pairs):
        self.pairs = sorted(tuple(p) for p in pairs)
        detail = ", ".join("%s -> %s" % (t, d) for t, d in self.pairs)
        super().__init__("unknown dependencies: " + detail)


class CycleError(SpecError):
    """The dependency graph is not a DAG (exit code 2)."""

    def __init__(self, cycle):
        self.cycle = list(cycle)
        super().__init__("dependency cycle: " + " -> ".join(self.cycle + [self.cycle[0]]))


# ---------------------------------------------------------------- spec parsing


def parse_spec(obj):
    """Validate a decoded spec.  Returns (tasks, workers)."""
    if not isinstance(obj, dict):
        raise SpecError("spec must be a JSON object")
    raw_tasks = obj.get("tasks")
    if not isinstance(raw_tasks, dict) or not raw_tasks:
        raise SpecError("spec.tasks must be a non-empty object")

    tasks = {}
    for name, body in raw_tasks.items():
        if not isinstance(name, str) or not name:
            raise SpecError("task names must be non-empty strings")
        if body is None:
            body = {}
        if not isinstance(body, dict):
            raise SpecError("task %r: body must be an object" % (name,))

        raw_deps = body.get("deps", [])
        if not isinstance(raw_deps, list):
            raise SpecError("task %r: deps must be a list" % (name,))
        deps = []
        for d in raw_deps:
            if not isinstance(d, str) or not d:
                raise SpecError("task %r: dependency names must be non-empty strings" % (name,))
            if d not in deps:
                deps.append(d)

        est = body.get("est", 1.0)
        if isinstance(est, bool) or not isinstance(est, (int, float)):
            raise SpecError("task %r: est must be a number" % (name,))
        est = float(est)
        if not (est > 0.0) or est == INF:
            raise SpecError("task %r: est must be positive and finite" % (name,))
        tasks[name] = {"deps": deps, "est": est}

    workers = obj.get("workers", 1)
    if isinstance(workers, bool) or not isinstance(workers, int) or workers < 1:
        raise SpecError("workers must be a positive integer")
    return tasks, workers


def build_graph(tasks):
    """Return (deps, rdeps) with sorted, deduplicated adjacency."""
    names = sorted(tasks)
    deps = {n: tuple(sorted(set(tasks[n]["deps"]))) for n in names}
    unknown = [(n, d) for n in names for d in deps[n] if d not in tasks]
    if unknown:
        raise UnknownDepError(unknown)
    rdeps = {n: [] for n in names}
    for n in names:
        for d in deps[n]:
            rdeps[d].append(n)
    for n in names:
        rdeps[n].sort()
    return deps, rdeps


# ------------------------------------------------------------------- cycle find


def find_cycle(deps):
    """Return one cycle as a list of names in execution order, or None.

    Traverses the E = {(d, n)} direction, so the returned list [x, y, ...]
    reads "x before y before ... before x".
    """
    rdeps = {n: [m for m in deps if n in deps[m]] for n in deps}
    for n in rdeps:
        rdeps[n].sort()

    WHITE, GREY, BLACK = 0, 1, 2
    color = {n: WHITE for n in deps}
    for root in sorted(deps):
        if color[root] != WHITE:
            continue
        color[root] = GREY
        path = [root]
        stack = [(root, iter(rdeps[root]))]
        while stack:
            node, it = stack[-1]
            descended = False
            for nxt in it:
                if color[nxt] == GREY:
                    return path[path.index(nxt):]
                if color[nxt] == WHITE:
                    color[nxt] = GREY
                    path.append(nxt)
                    stack.append((nxt, iter(rdeps[nxt])))
                    descended = True
                    break
            if not descended:
                color[node] = BLACK
                stack.pop()
                path.pop()
    return None


# ------------------------------------------------------------------- orderings


def _kahn(deps, rdeps):
    """Kahn's algorithm with a min-heap of ready tasks.  Raises CycleError."""
    indeg = {n: len(deps[n]) for n in deps}
    ready = [n for n in deps if indeg[n] == 0]
    heapq.heapify(ready)
    order = []
    while ready:
        n = heapq.heappop(ready)
        order.append(n)
        for m in rdeps[n]:
            indeg[m] -= 1
            if indeg[m] == 0:
                heapq.heappush(ready, m)
    if len(order) != len(deps):
        raise CycleError(find_cycle(deps) or [])
    return order


def topological_order(deps):
    """A serial build order; lexicographically smallest ready task first."""
    _, rdeps = build_graph({n: {"deps": list(deps[n]), "est": 1.0} for n in deps})
    return _kahn(deps, rdeps)


def waves(deps):
    """Parallel build waves.  wave k holds every task whose longest
    dependency chain has k edges.  Each wave is sorted."""
    _, rdeps = build_graph({n: {"deps": list(deps[n]), "est": 1.0} for n in deps})
    order = _kahn(deps, rdeps)
    level = {n: 0 for n in deps}
    for n in order:
        if deps[n]:
            level[n] = 1 + max(level[d] for d in deps[n])
    depth = max(level.values()) + 1
    out = [[] for _ in range(depth)]
    for n in sorted(level):
        out[level[n]].append(n)
    return out


def critical_path(deps, est):
    """Longest weighted path.  Returns (length, path).

    Ties (equal-length predecessors / equal-length ends) resolve to the
    lexicographically smallest name, so the answer is unique.
    """
    _, rdeps = build_graph({n: {"deps": list(deps[n]), "est": float(est[n])} for n in deps})
    order = _kahn(deps, rdeps)
    dist = {}
    pred = {}
    for n in order:
        best, best_d = 0.0, None
        for d in deps[n]:
            if best_d is None or dist[d] > best or (dist[d] == best and d < best_d):
                best, best_d = dist[d], d
        dist[n] = best + est[n]
        pred[n] = best_d
    if not dist:
        return 0.0, []
    longest = max(dist.values())
    end = min(n for n in dist if dist[n] == longest)
    path = []
    cur = end
    while cur is not None:
        path.append(cur)
        cur = pred[cur]
    path.reverse()
    return dist[end], path


# ------------------------------------------------------------------- scheduling


def schedule(deps, est, workers):
    """Greedy list scheduling on `workers` identical machines.

    Priority: lexicographically smallest ready task.  A worker takes the
    first ready task at or after the moment it becomes free.  Returns a list
    of dicts sorted by (start, name).
    """
    _, rdeps = build_graph({n: {"deps": list(deps[n]), "est": float(est[n])} for n in deps})
    order = _kahn(deps, rdeps)
    if not order:
        return []

    indeg = {n: len(deps[n]) for n in deps}
    ready = [n for n in deps if indeg[n] == 0]
    heapq.heapify(ready)
    free = [(0.0, w) for w in range(workers)]
    heapq.heapify(free)
    running = []  # (end, name, worker)
    out = []
    now = 0.0
    placed = 0

    while placed < len(order):
        while running and running[0][0] <= now:
            end, name, w = heapq.heappop(running)
            heapq.heappush(free, (end, w))
            for m in rdeps[name]:
                indeg[m] -= 1
                if indeg[m] == 0:
                    heapq.heappush(ready, m)
        while ready and free and free[0][0] <= now:
            _, w = heapq.heappop(free)
            name = heapq.heappop(ready)
            end = now + est[name]
            out.append({"task": name, "start": now, "end": end, "worker": w})
            heapq.heappush(running, (end, name, w))
            placed += 1
        if placed == len(order):
            break
        # Every completion at or before `now` has been drained, and workers are
        # only released by those completions, so the next event is strictly in
        # the future.  A DAG always leaves a runnable task here; the guard turns
        # any hypothetical scheduler bug into a loud failure instead of a spin.
        nxt = running[0][0] if running else None
        if nxt is None or nxt <= now:
            raise RuntimeError("scheduler stalled with %d task(s) unplaced" % (len(order) - placed))
        now = nxt

    out.sort(key=lambda r: (r["start"], r["task"]))
    return out


# ------------------------------------------------------------------------ plan


def plan(obj):
    """Full plan dict for a decoded spec."""
    tasks, workers = parse_spec(obj)
    deps, _ = build_graph(tasks)
    est = {n: tasks[n]["est"] for n in deps}
    order = topological_order(deps)
    wv = waves(deps)
    cp_len, cp_path = critical_path(deps, est)
    sched = schedule(deps, est, workers)
    return {
        "tasks": sorted(deps),
        "order": order,
        "waves": wv,
        "critical_path": {"tasks": cp_path, "length": cp_len},
        "workers": workers,
        "makespan": max((r["end"] for r in sched), default=0.0),
        "serial_time": sum(est.values()),
        "schedule": sched,
    }


# ------------------------------------------------------------------------- CLI


def main(argv=None):
    argv = list(sys.argv[1:] if argv is None else argv)
    if not argv or argv[0] in ("-h", "--help"):
        sys.stderr.write("usage: deporder.py plan <spec.json> [--workers N]\n")
        return 0 if argv else 4
    if argv[0] != "plan":
        sys.stderr.write("unknown command %r\n" % (argv[0],))
        return 4

    rest = argv[1:]
    workers = None
    path = None
    i = 0
    while i < len(rest):
        if rest[i] == "--workers":
            if i + 1 >= len(rest):
                sys.stderr.write("--workers needs a value\n")
                return 4
            try:
                workers = int(rest[i + 1])
            except ValueError:
                sys.stderr.write("--workers must be an integer\n")
                return 4
            i += 2
        elif path is None:
            path = rest[i]
            i += 1
        else:
            sys.stderr.write("unexpected argument %r\n" % (rest[i],))
            return 4

    if path is None:
        sys.stderr.write("missing spec path\n")
        return 4
    try:
        with open(path, "r", encoding="utf-8") as fh:
            obj = json.load(fh)
    except OSError as exc:
        sys.stderr.write("cannot read %s: %s\n" % (path, exc))
        return 4
    except ValueError as exc:
        sys.stderr.write("invalid JSON in %s: %s\n" % (path, exc))
        return 4

    if workers is not None:
        if not isinstance(obj, dict):
            sys.stderr.write("spec must be a JSON object\n")
            return 4
        obj = dict(obj)
        obj["workers"] = workers

    try:
        result = plan(obj)
    except UnknownDepError as exc:
        sys.stderr.write(str(exc) + "\n")
        return 3
    except CycleError as exc:
        sys.stderr.write(str(exc) + "\n")
        return 2
    except SpecError as exc:
        sys.stderr.write(str(exc) + "\n")
        return 4

    json.dump(result, sys.stdout, indent=2, sort_keys=True)
    sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())