from typing import Set
from api.task import Task, Status
from api.execution_context import ExecutionContext


class Scheduler:
    """
    DAG-based execution scheduler.
    Decides WHEN a Task can run based on port readiness and orchestrates execution.
    """

    def __init__(self):
        self._completed: Set[Task] = set()
        self._failed: Set[Task] = set()

    def run(self, root: Task, context: ExecutionContext):
        self._execute_task(root, context)

    def _execute_task(self, task: Task, context: ExecutionContext):
        if task.status in (Status.COMPLETED, Status.FAILED, Status.STOPPED):
            return

        if hasattr(task, "children"):
            self._run_workflow(task, context)
            return

        if not task.check_inputs_ready(context):
            print(f"[Scheduler] Task '{task.name}' not ready, skipping for now.")
            return

        try:
            task._set_status(Status.RUNNING)
            task.run(context)
            task._set_status(Status.COMPLETED)
            self._completed.add(task)
        except Exception as e:
            task._set_status(Status.FAILED)
            self._failed.add(task)
            print(f"[Scheduler] Task '{task.name}' failed: {e}")
            raise

    def _run_workflow(self, workflow: Task, context: ExecutionContext):
        workflow._set_status(Status.RUNNING)

        remaining = set(workflow.children)

        while remaining:
            progress_made = False

            for child in list(remaining):
                if child.check_inputs_ready(context):
                    self._execute_task(child, context)
                    remaining.remove(child)
                    progress_made = True

            if not progress_made:
                raise RuntimeError(
                    f"Workflow '{workflow.name}' deadlock: unresolved dependencies"
                )

        if any(c.status == Status.FAILED for c in workflow.children):
            workflow._set_status(Status.FAILED)
        else:
            workflow._set_status(Status.COMPLETED)