import subprocess
from collections import deque
from concurrent.futures import Future, ProcessPoolExecutor

from brake.graph import TargetGraph
from brake.model import Target


class TargetGraphRunner:
    def __init__(self, target_graph: TargetGraph):
        self.target_graph = target_graph
        self.pool: ProcessPoolExecutor | None = None

    def submit_target(self, target_name: str) -> Future:
        """Submit a target by name to the worker pool"""
        if self.pool:
            return self.pool.submit(self.run_target, self.target_graph[target_name])
        raise RuntimeError("The process pool has not yet been initialized")

    @staticmethod
    def run_target(target: Target):
        for cmd in target.commands:
            print(f"[{target.type}:{target.name}] {cmd}")
            subprocess.run(cmd, shell=True, check=True)

    @staticmethod
    def wait_for_one_command_to_complete(commands: dict[Future, str]) -> str | None:
        try:
            done = next(f for f in commands if f.done())
        except StopIteration:
            return
        else:
            return commands.pop(done)

    def run(self, target_name: str | None, max_workers: int):
        if not target_name:
            if default_target := self.target_graph.default_target:
                target_name = default_target.name
            else:
                raise ValueError("A target is required to run as the graph has no default target")

        dependents = self.target_graph.target_dependency_graph(target_name)
        remaining_deps = self.target_graph.target_subgraph_dependency_counters(target_name)

        # targets ready to run
        ready = deque(name for name, n in remaining_deps.items() if n == 0)

        running, completed = {}, set()

        with ProcessPoolExecutor(max_workers=max_workers) as self.pool:
            while ready or running:
                # schedule new targets
                while ready and len(running) < max_workers:
                    name = ready.popleft()
                    future = self.submit_target(name)
                    running[future] = name

                if completed_target_name := self.wait_for_one_command_to_complete(running):
                    completed.add(completed_target_name)

                    # unlock dependents
                    for dep in dependents[completed_target_name]:
                        remaining_deps[dep] -= 1
                        if remaining_deps[dep] == 0:
                            ready.append(dep)
