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

from brake.colors import Color, colorize
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, dry_run: bool = False, color: bool = True) -> 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], dry_run, color)
        raise RuntimeError("The process pool has not yet been initialized")

    @staticmethod
    def run_target(target: Target, dry_run: bool = False, color: bool = True):
        for cmd in target.commands:
            task = f"[{target.type}:{target.name}] {cmd}"
            output_color = Color.from_str(target.name)
            if not color:
                print(task)
            else:
                print(colorize(task, output_color))
            if not dry_run:
                if not color:
                    subprocess.run(cmd, shell=True, check=True, encoding="utf-8")
                else:
                    process = subprocess.run(
                        cmd, shell=True, check=True, capture_output=True, encoding="utf-8"
                    )
                    if not process.stdout:
                        continue
                    print(colorize(process.stdout, output_color), end="")

    @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 None
        else:
            return commands.pop(done)

    def run(
        self, target_name: str | None, max_workers: int, dry_run: bool = False, color: bool = True
    ) -> None:
        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: dict[Future, str] = {}
        completed: set[str] = set()

        # don't colorize the output if there's only a single task to run anyway
        if len(remaining_deps) == 1 and color:
            color = False

        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, dry_run=dry_run, color=color)
                    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)
