# Copyright 2026 Q-CTRL. All rights reserved.
#
# Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
# copying or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. You may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#    https://q-ctrl.com/terms
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS. See the
# License for the specific language.

from __future__ import annotations

__all__ = ("Runtime",)

import contextlib
import graphlib
import warnings
from dataclasses import dataclass
from typing import TYPE_CHECKING, Self

from boulderopalscaleup.runtime._sync.executable import SyncExecIntermediate
from boulderopalscaleup.runtime.abstract import (
    ExecutableIntermediate,
    ExecutableResult,
    ExecutionMode,
    RuntimeOptions,
    System,
    SystemInfo,
)
from boulderopalscaleup.runtime.manifest import ManifestGraph

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator

    from boulderopalscaleupsdk.runtime.manifest import Manifest, ManifestNode


class Runtime:
    """
    Manages hardware systems and executes compiled programs.

    A runtime holds a set of registered hardware systems and executes compiled
    programs against them. It supports the context manager protocol for automatic
    resource cleanup.

    Use ``Runtime.new()`` as the preferred constructor.

    Attributes
    ----------
    is_closed : bool
        ``True`` once the runtime and all its systems have been closed.
    systems : dict[str, System]
        Registered hardware systems, keyed by system identifier.
    """

    def __init__(
        self,
        options: RuntimeOptions,
    ) -> None:
        self._closed: bool = False
        self._close_on_error: bool = options.close_on_system_error
        self._systems: dict[str, System] = {}
        if not isinstance(options.exec_mode, ExecutionMode):  # pragma: no cover
            # Don't rely on type hinting for public API.
            # We are stricter here, because there is a possibility that we build new systems from
            # the Runtime.add_system method whilst this Runtime is in an erroneous state, leading
            # to poor resource management edge-cases.
            raise ValueError(f"Unsupported execution mode '{self._exec_mode!r}'")  # noqa: TRY004
        self._exec_mode = options.exec_mode
        self._options = options

    @property
    def is_closed(self) -> bool:
        return self._closed

    @property
    def systems(self) -> dict[str, System]:
        return self._systems

    @classmethod
    def new(cls, options: RuntimeOptions | None = None) -> Self:
        """Create a new runtime with the provided options.

        Parameters
        ----------
        options : RuntimeOptions, optional
            Configuration for the runtime. Defaults to ``RuntimeOptions()`` if not provided.

        Returns
        -------
        Runtime
            A new runtime instance.
        """
        return cls(options or RuntimeOptions())

    def __enter__(self) -> Self:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        self.close()

    def close(self) -> None:
        """Close the runtime and all registered systems.

        Attempts to close every system even if one or more raise exceptions,
        collecting all errors and re-raising them as an ``ExceptionGroup``.

        Raises
        ------
        ExceptionGroup
            If one or more systems fail to close.
        """
        errors: dict[str, Exception] = {}
        for sys in self._systems.values():
            try:  # Always try close as many systems as possible!
                sys.close()
            except Exception as exc:  # noqa: BLE001
                errors[sys.implements] = exc

        # Release references and close runtime before raising in case caller wraps this call.
        self._systems = {}
        self._closed = True
        if errors:
            failed_sys_str = ", ".join(errors.keys())
            raise ExceptionGroup(
                f"Runtime failed to close {len(errors)} systems ({failed_sys_str}).",
                list(errors.values()),
            )

    @contextlib.contextmanager
    def _wrap_close(self) -> Iterator[None]:
        """Used to ensure we honour close_on_error runtime option."""
        try:
            yield
        except Exception:
            if bool(self._close_on_error):
                self.close()
            raise

    def add_system(self, system: System | SystemInfo) -> None:
        """Register a system with this runtime.

        Accepts either a ``System`` instance directly or a ``SystemInfo`` description
        object; in the latter case, the system is built via ``SystemInfo.build_system()``.

        Parameters
        ----------
        system : System | SystemInfo
            The system or system-information object to register.

        Raises
        ------
        RuntimeError
            If the runtime is already closed.
        ValueError
            If the system's execution mode does not match the runtime's execution mode.
        """
        if self._closed:
            raise RuntimeError("Runtime is closed, cannot add system.")
        with self._wrap_close():
            match system:
                case System():  # pragma: no cover
                    if system.exec_mode != self._exec_mode:
                        raise ValueError(
                            "Runtime does not support multiple execution modes. This runtime is "
                            f"set with execution mode '{self._exec_mode}' but got a system with "
                            f"execution mode '{system.exec_mode}'.",
                        )

                    if system.implements in self._systems:  # pragma: no cover
                        # TODO: Allow for multiple systems instances of same implementation.
                        #       See https://qctrl.atlassian.net/browse/SCUP-3199
                        warnings.warn(
                            "Runtime already has added a system implementing "
                            f"'{system.implements}', overriding...",
                            stacklevel=2,
                        )
                    self._systems[system.implements] = system
                case SystemInfo():
                    self.add_system(system.build_system())
                case _:  # pragma: no cover
                    raise ValueError(
                        f"Unsupported type '{type(system).__name__}' for system, expected System or "  # noqa: E501
                        "SystemInfo.",
                    )

    def bind(self, manifest: Manifest) -> _Executable:
        """Compile a manifest into an executable against registered systems.

        Parameters
        ----------
        manifest : Manifest
            The compiled manifest to execute.

        Returns
        -------
        Executable
            An executable ready to be run via its ``run()`` method.

        Raises
        ------
        RuntimeError
            If the runtime is already closed.
        ValueError
            If the manifest requires a system that has not been registered.
        """
        if self._closed:
            raise RuntimeError("Runtime is closed, cannot bind manifest.")
        with self._wrap_close():
            # TODO: Check for cycle error. See https://qctrl.atlassian.net/browse/SCUP-3195
            graph = ManifestGraph.from_manifest(manifest)
            intermediate: ExecutableIntermediate
            match self._exec_mode:
                case ExecutionMode.SYNC:
                    intermediate = SyncExecIntermediate.init_from_manifest(manifest)

            # Allow systems to prepare resources before node binding
            for system in self._systems.values():
                system.prepare_resources(manifest, graph)

            for node in _iter_graph_ordered(graph):
                if node.system not in self._systems:
                    raise ValueError(
                        f"Node {node.id} requires system of type '{node.system}' but none found.",
                    )
                system = self._systems[node.system]

                # TODO: Check system is not closed.
                #       See https://qctrl.atlassian.net/browse/SCUP-3198
                intermediate = system.bind(node, graph, intermediate)

            # TODO: Add observability to execution. See https://qctrl.atlassian.net/browse/SCUP-3200
            return _Executable(
                _runtime=self,
                _runtime_options=self._options,
                __executable_intermediate__=intermediate,
            )

    def run(self, manifest: Manifest) -> ExecutableResult:
        """Bind and immediately execute a manifest.

        Shorthand for ``runtime.bind(manifest).run()``.

        Parameters
        ----------
        manifest : Manifest
            The compiled manifest to execute.

        Returns
        -------
        ExecutableResult
            The result of execution, containing output data and timing metadata.

        Raises
        ------
        RuntimeError
            If the runtime is already closed.
        ValueError
            If the manifest requires a system that has not been registered.
        """
        return self.bind(manifest).run()


@dataclass
class _Executable:
    _runtime: Runtime | None
    _runtime_options: RuntimeOptions
    __executable_intermediate__: ExecutableIntermediate

    def run(self) -> ExecutableResult:
        if self._runtime is None or self._runtime.is_closed:  # pragma: no cover
            # Once we've closed the parent runtime (see try-catch block below), we don't want to
            # allow anymore retries; even if we did try, the run method should still run into errors
            # provided the downstream systems have guarded their respective execution unit.
            # Despite this, we manually check here as to not rely on downstream system
            # implementations.
            raise RuntimeError("This executable is not bound to an open runtime.")
        try:
            return self.__executable_intermediate__.run()
        except Exception:
            if self._runtime_options.close_on_exec_error:
                self._runtime.close()
                self._runtime = None
            raise

        # NOTE: this executable can be retried, hence we don't release the parent runtime reference
        #       unless explicitly configured to do so!


def _iter_graph_ordered(graph: ManifestGraph) -> Iterable[ManifestNode]:
    sorter = graphlib.TopologicalSorter(graph.predecessors)
    for node_name in sorter.static_order():
        yield graph.nodes[node_name]
