execution

class ExecutionManager:
    """ Manages workflow executions and their event queues.
    """

    def __init__(self):
        self._lock = threading.Lock()
        self._next_run_id = 0
        self._runs: Dict[int, Dict[str, Any]] = {}

    def __new_run(self) -> Tuple[int, Dict[str, Any]]:
        with self._lock:
            self._next_run_id += 1
            run = {
                "status": "running",
                "cancel_event": threading.Event(),
                "queue": RunQueue(),
            }
            self._runs[self._next_run_id] = run
            return self._next_run_id, run

    def __get_run_by_id(self, run_id: int) -> Optional[Dict[str, Any]]:
        with self._lock:
            return self._runs.get(run_id)

    def _set_status(
        self,
        run_id: int,
        status: str,
        queue: RunQueue,
    ):
        with self._lock:
            run = self._runs.get(run_id)
            if run:
                run["status"] = status

        _STATUS_TO_EVENT = {
            "running": "execution.status",
            "done": "execution.completed",
            "failed": "execution.failed",
            "cancelled": "execution.failed",
        }
        event_type = _STATUS_TO_EVENT.get(status, "execution.status")
        queue.put({
            "event_type": event_type,
            "run_id": run_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "execution_status": status,
        })

    def get_status(self, run_id: int) -> Dict[str, Any]:
        run = self.__get_run_by_id(run_id)
        return {
            "run_id": run_id,
            "status": run["status"] if run else "not_found",
        }

    def get_queue(self, run_id: int) -> Optional[RunQueue]:
        run = self.__get_run_by_id(run_id)
        return run["queue"] if run else None

    def start_execution(self, workflow: Workflow) -> Dict[str, Any]:
        """ Start a new local execution of the provided workflow.
        """
        run_id, run = self.__new_run()
        run["queue"].put({
            "event_type": "execution.status",
            "run_id": run_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "execution_status": "running",
        })

        # perform the execution in a new thread
        threading.Thread(
            target=_run_workflow_with_monitoring,
            args=(
                workflow,
                run_id,
                run["queue"],
                run["cancel_event"],
                self._set_status,
            ),
            name=f"ayon-workflow-run-{run_id}",
            daemon=True,
        ).start()
        return self.get_status(run_id)

    def start_submit(
        self,
        workflow: Workflow,
        backend_dir: str,
        dispatch_graph_name: Optional[str] = None,
        project: Optional[str] = None,
    ) -> Dict[str, Any]:
        """ Export workflow, then submit to farm and watch backend_dir.
        """
        import tempfile
        from contextlib import suppress
        from ayon_workflow.workflow_execution.renderfarm import (
            submit_workflow_to_farm,
        )

        run_id, run = self.__new_run()
        queue = run["queue"]

        temp_path = None
        try:
            # Export workflow to temporary file
            with tempfile.NamedTemporaryFile(
                dir=backend_dir,
                suffix="_temp.json",
                delete=False
            ) as fh:
                temp_path = fh.name
            workflow.export_to_file(temp_path)

            # Submit workflow to farm
            submit_workflow_to_farm(
                temp_path,
                backend_dir=backend_dir,
                dispatch_graph_name=dispatch_graph_name,
                project=project,
            )

        except Exception as error:
            error_message = str(error)
            with self._lock:
                run["status"] = "failed"
            queue.put({
                "event_type": "execution.log",
                "run_id": run_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "level": "ERROR",
                "message": error_message,
            })
            queue.put({
                "event_type": "execution.failed",
                "run_id": run_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "message": error_message,
            })
            queue.put({
                "event_type": "execution.end",
                "run_id": run_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            })
            queue.close()
            raise
        finally:
            if temp_path:
                with suppress(OSError):
                    os.unlink(temp_path)

        queue.put({
            "event_type": "execution.status",
            "run_id": run_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "execution_status": "running",
        })

        # Start watching the backend in a separate thread.
        watcher = _BackendDirWatcher(backend_dir, run_id, queue)
        threading.Thread(
            target=_watch_submit,
            args=(watcher, run_id, queue, self._set_status),
            name=f"ayon-workflow-watch-{run_id}",
            daemon=True,
        ).start()

        return self.get_status(run_id)

    def cancel(self, run_id: int) -> bool:
        """ Cancel a running workflow by its id.
        """
        with self._lock:
            run = self._runs.get(run_id)
            if run is None or run["status"] != "running":
                return False

            # TODO: current implementation is a cosmetic cancel,
            # the workflow will continue to run. Need to implement a proper
            # cancellation mechanism for execution and submission.
            run["cancel_event"].set()
            run["status"] = "cancelled"

        run["queue"].put({
            "event_type": "execution.failed",
            "run_id": run_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "execution_status": "cancelled",
        })
        return True


    def cancel_all_executions(self):
        """ Cancel all currently running executions.
        """
        with self._lock:
            run_ids = [
                run_id
                for run_id, run in self._runs.items()
                if run["status"] == "running"
            ]
        for run_id in run_ids:
            self.cancel(run_id)

class RunQueue:
    """ Per run_id event queue with index/cursor-based replay.
    """

    def __init__(self):
        self._lock = threading.Condition(threading.Lock())
        self._events: List[dict] = []
        self._closed: bool = False

    def put(self, event: dict):
        with self._lock:
            self._events.append(event)
            self._lock.notify_all()  # notify all waiting threads

    def close(self):
        with self._lock:
            self._closed = True
            self._lock.notify_all()  # notify all waiting threads

    def get_from(
        self,
        idx: int,
        timeout: float = 0.5,
    ) -> Tuple[int, List[dict], bool]:
        """ Block until new events are available or timeout.
        """
        with self._lock:
            self._lock.wait_for(
                lambda: len(self._events) > idx or self._closed,
                timeout=timeout,  # timeout while no new events
            )
            events = self._events[idx:]
            new_idx = idx + len(events)
            return new_idx, events, self._closed