"""DAG (Directed Acyclic Graph) for Boris milestone dependency management.

Provides topological sort (Kahn's algorithm), cycle detection, critical path
calculation, parallel batch computation, and dependency validation.
"""
from __future__ import annotations

from collections import deque
from state import Milestone, Plan


class DependencyError(Exception):
    """Raised when the dependency graph is invalid (cycles, missing refs)."""


class DependencyGraph:
    """Directed Acyclic Graph for milestone dependencies.

    Build from a Plan, then query for execution order, parallel batches,
    or ready milestones.
    """

    def __init__(self, milestones: list[Milestone]):
        self._milestones = {m.id: m for m in milestones}
        self._adj: dict[str, list[str]] = {m.id: [] for m in milestones}
        self._in_degree: dict[str, int] = {m.id: 0 for m in milestones}

        for m in milestones:
            for dep in m.depends_on:
                if dep in self._adj:  # skip unknown deps (caught by validate)
                    self._adj[dep].append(m.id)
                    self._in_degree[m.id] += 1

    # --- Validation (call at plan creation time) ---

    def validate(self) -> None:
        """Validate the graph: check for missing deps and cycles.

        Raises DependencyError with a clear message on failure.
        """
        self._check_missing_deps()
        self._check_cycles()

    def _check_missing_deps(self) -> None:
        """Ensure all depends_on IDs reference milestones that exist."""
        valid_ids = set(self._milestones.keys())
        for m in self._milestones.values():
            bad = [d for d in m.depends_on if d not in valid_ids]
            if bad:
                raise DependencyError(
                    f"Milestone {m.id} depends on unknown milestone(s): {', '.join(bad)}. "
                    f"Valid IDs: {', '.join(sorted(valid_ids))}"
                )

    def _check_cycles(self) -> None:
        """Detect cycles using Kahn's algorithm.

        If not all nodes are consumed, there's a cycle.
        """
        in_deg = dict(self._in_degree)
        queue = deque(nid for nid, d in in_deg.items() if d == 0)
        visited = 0

        while queue:
            node = queue.popleft()
            visited += 1
            for neighbor in self._adj[node]:
                in_deg[neighbor] -= 1
                if in_deg[neighbor] == 0:
                    queue.append(neighbor)

        if visited != len(self._milestones):
            remaining = [nid for nid, d in in_deg.items() if d > 0]
            raise DependencyError(
                f"Circular dependency detected among milestones: {', '.join(remaining)}. "
                f"Check depends_on fields for cycles."
            )

    # --- Topological sort (Kahn's algorithm) ---

    def topological_sort(self) -> list[str]:
        """Return milestone IDs in a valid topological order.

        Uses Kahn's algorithm. Ties are broken by milestone ID for
        deterministic ordering.
        """
        in_deg = dict(self._in_degree)
        # Use a sorted list to break ties deterministically
        queue = sorted(nid for nid, d in in_deg.items() if d == 0)
        result = []

        while queue:
            node = queue.pop(0)
            result.append(node)
            new_ready = []
            for neighbor in self._adj[node]:
                in_deg[neighbor] -= 1
                if in_deg[neighbor] == 0:
                    new_ready.append(neighbor)
            # Insert new ready nodes in sorted order
            queue.extend(new_ready)
            queue.sort()

        return result

    # --- Critical path ---

    def critical_path_lengths(self) -> dict[str, int]:
        """Compute the longest path from each node to any sink.

        Returns a dict mapping milestone ID -> length of longest path
        starting from that milestone. Milestones on the critical path
        have the highest values.
        """
        order = self.topological_sort()
        # Process in reverse topological order
        longest = {nid: 0 for nid in self._milestones}
        for nid in reversed(order):
            for neighbor in self._adj[nid]:
                longest[nid] = max(longest[nid], 1 + longest[neighbor])
        return longest

    # --- Parallel batches (topological layers) ---

    def parallel_batches(self, done_ids: set[str] | None = None,
                         skip_ids: set[str] | None = None) -> list[list[str]]:
        """Return batches of milestones that can run concurrently.

        Each batch is a topological layer: all milestones in a batch have
        their dependencies satisfied by previous batches. Within each batch,
        milestones are sorted by critical path length (longest first) for
        priority, falling back to ID for determinism.

        Args:
            done_ids: Milestones already completed/skipped (excluded from output).
            skip_ids: Milestones to exclude entirely.
        """
        done = done_ids or set()
        skip = skip_ids or set()
        excluded = done | skip

        # Build filtered in-degree map
        in_deg = {}
        for nid in self._milestones:
            if nid in excluded:
                continue
            deg = sum(1 for d in self._milestones[nid].depends_on
                      if d not in excluded)
            in_deg[nid] = deg

        crit = self.critical_path_lengths()
        batches = []

        while in_deg:
            # Collect all nodes with in-degree 0
            batch = [nid for nid, d in in_deg.items() if d == 0]
            if not batch:
                # Remaining nodes form a cycle (shouldn't happen after validate)
                break
            # Sort: longest critical path first, then by ID
            batch.sort(key=lambda nid: (-crit.get(nid, 0), nid))
            batches.append(batch)

            # Remove batch nodes and update in-degrees
            for nid in batch:
                del in_deg[nid]
                for neighbor in self._adj[nid]:
                    if neighbor in in_deg:
                        in_deg[neighbor] -= 1

        return batches

    # --- Ready milestones (for runtime use) ---

    def get_ready(self, done_ids: set[str]) -> list[str]:
        """Return IDs of pending milestones whose deps are all in done_ids.

        Results are sorted by critical path (longest first) so the most
        important milestones start first.
        """
        crit = self.critical_path_lengths()
        ready = []
        for nid, m in self._milestones.items():
            if nid in done_ids:
                continue
            if all(d in done_ids for d in m.depends_on):
                ready.append(nid)
        ready.sort(key=lambda nid: (-crit.get(nid, 0), nid))
        return ready

    # --- Execution order (for sequential mode) ---

    def execution_order(self) -> list[str]:
        """Return the optimal milestone execution order for sequential mode.

        Uses topological sort with critical path priority: at each step,
        the milestone with the longest remaining path is chosen first.
        This ensures dependencies on the critical path are resolved ASAP.
        """
        crit = self.critical_path_lengths()
        in_deg = dict(self._in_degree)
        # Seed with all zero-indegree nodes, sorted by critical path desc
        queue = sorted(
            (nid for nid, d in in_deg.items() if d == 0),
            key=lambda nid: (-crit.get(nid, 0), nid),
        )
        result = []

        while queue:
            node = queue.pop(0)
            result.append(node)
            new_ready = []
            for neighbor in self._adj[node]:
                in_deg[neighbor] -= 1
                if in_deg[neighbor] == 0:
                    new_ready.append(neighbor)
            queue.extend(new_ready)
            queue.sort(key=lambda nid: (-crit.get(nid, 0), nid))

        return result


# --- Module-level convenience functions ---

def validate_plan_dependencies(plan: Plan) -> None:
    """Validate a plan's dependency graph at creation time.

    Raises DependencyError if there are cycles or missing references.
    """
    graph = DependencyGraph(plan.milestones)
    graph.validate()


def plan_execution_order(plan: Plan) -> list[str]:
    """Return optimal sequential execution order for a plan's milestones."""
    graph = DependencyGraph(plan.milestones)
    return graph.execution_order()


def get_parallel_batches(plan: Plan,
                         done_ids: set[str] | None = None,
                         skip_ids: set[str] | None = None) -> list[list[str]]:
    """Return batches of milestones that can run concurrently."""
    graph = DependencyGraph(plan.milestones)
    return graph.parallel_batches(done_ids=done_ids, skip_ids=skip_ids)
