from typing import Set
from radiuma_api.task import Task, Status
from radiuma_api.execution_context import ExecutionContext
# Write Codes Here:


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):
        """
        Entry point: execute a Workflow or Task graph.
        """
        self._execute_task(root, context)

    def _execute_task(self, task: Task, context: ExecutionContext):
        """
        Execute a task respecting DAG dependencies.
        If 'task' is a Workflow (has children), delegate to _run_workflow.
        Otherwise, run the leaf task when inputs are ready.
        """
        # Skip finished tasks
        if task.status in (Status.COMPLETED, Status.FAILED, Status.STOPPED):
            return

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

        # Leaf Task: must have inputs ready
        if not task.check_inputs_ready(context):
            # Optional: diagnostic log
            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)
            # Optional: diagnostic log
            print(f"[Scheduler] Task '{task.name}' failed: {e}")
            # Re-raise to surface error if caller needs to handle
            raise

    def _run_workflow(self, workflow: Task, context: ExecutionContext):
        """
        Execute children respecting dependencies.
        Progress until all children are executed or a deadlock is detected.
        """
        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:
                # Deadlock: dependencies not resolvable with current data
                raise RuntimeError(
                    f"Workflow '{workflow.name}' deadlock: unresolved dependencies"
                )

        # Final workflow status based on children results
        if any(c.status == Status.FAILED for c in workflow.children):
            workflow._set_status(Status.FAILED)
        else:
            workflow._set_status(Status.COMPLETED)