# 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.

"""
Abstract type definitions for Boulder Opal Scale Up runtime target.

This module provides abstract definitions of the runtime the Boulder Opal Scale Up compiler will
target.

Basic usage
-----------

```py
manifest = compile(...)  # Boulder Opal Scale Up compiler

runtime = Runtime()
runtime.add_system(ASystemInfoDescription())
executable = runtime.bind(manifest)

results = executable()
```

Overview
--------

Systems provide functionality --we expect primarily through interfacing with external control
hardware-- which the compiler can target by producing the corresponding ManifestNode.

A Runtime manages a set of systems and binds a Manifest (a collection of constituent manifest nodes)
to the corresponding system functions. The binding process is a process in which the Runtime and its
component Systems:

- Map manifest nodes into self-contained execution units
- Collect and order the units under a single executable
- Ensure units can retrieve results of dependent units (via result proxies)
- Ensure units can access the required resources (via resource proxies)
"""

from __future__ import annotations

__all__ = (
    "CalibrationError",
    "Executable",
    "ExecutableIntermediate",
    "ExecutableResult",
    "ExecutionMetadata",
    "ExecutionMode",
    "ResourceProxy",
    "ResultProxy",
    "Runtime",
    "RuntimeOptions",
    "System",
    "SystemInfo",
)

import abc
import enum
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, Self, TypeVar, runtime_checkable

if TYPE_CHECKING:
    from contextlib import AbstractContextManager

    from boulderopalscaleupsdk.runtime.manifest import Manifest, ManifestNode, SupportedSystem

    from boulderopalscaleup.runtime.manifest import ManifestGraph


@dataclass
class ExecutionMetadata:
    """Metadata collected during execution.

    Attributes
    ----------
    duration : float
        Total execution time in milliseconds, summed across all timed nodes.
    """

    duration: float = 0.0


# TODO: Backwards compatible data structure for calibration error messages.
# Ideally this should be raised as an exception during execution.
# https://qctrl.atlassian.net/browse/SCUP-3848
@dataclass
class CalibrationError:
    """Calibration error message."""

    message: str


# TODO: Add documentation for expected exceptions raised.
#       See https://qctrl.atlassian.net/browse/SCUP-3196

_ResultT_co = TypeVar("_ResultT_co", covariant=True)
_ResT_co = TypeVar("_ResT_co", covariant=True)


class ResourceProxy(Protocol[_ResT_co]):
    """Provides managed access to resources for executables."""

    @abc.abstractmethod
    def get(self) -> AbstractContextManager[_ResT_co]:
        raise NotImplementedError


class ResultProxy(Protocol[_ResultT_co]):
    """Provides pull-based access between steps in a executable."""

    @abc.abstractmethod
    def get(self) -> _ResultT_co:
        raise NotImplementedError


class ExecutableResult(Protocol):
    """The result of an executable."""

    @property
    def output(self) -> Any:
        raise NotImplementedError

    @property
    def metadata(self) -> ExecutionMetadata:
        raise NotImplementedError

    def get_intermediate(self, node_id: str) -> Any:
        raise NotImplementedError


_UnitT_contra = TypeVar("_UnitT_contra", contravariant=True)


class ExecutableIntermediate(Protocol[_UnitT_contra]):
    """The intermediate executable object used in construction of the final executable."""

    def set_unit(self, node_id: str, unit: _UnitT_contra) -> None:
        """Set the execution unit in the executable for a node in the target manifest."""
        raise NotImplementedError

    def get_result_proxy(self, node_id: str) -> ResultProxy:
        """Get a result proxy matching a node in the target manifest."""
        raise NotImplementedError

    def run(self) -> ExecutableResult:
        """The run method of this executable."""
        raise NotImplementedError

    @classmethod
    @abc.abstractmethod
    def init_from_manifest(cls, manifest: Manifest) -> Self:
        raise NotImplementedError


_E = TypeVar("_E", bound=ExecutableIntermediate)


class Executable(Protocol[_E]):
    """
    The final executable produced from runtime binding.

    This object should wrap the executable intermediate object that was used to create this. This
    wrapper is needed to allow runtimes to modify execution behaviour (see RuntimeOptions).
    """

    __executable_intermediate__: _E

    @abc.abstractmethod
    def run(self) -> ExecutableResult:
        raise NotImplementedError


class ExecutionMode(enum.StrEnum):
    SYNC = "SYNC"


@runtime_checkable
class System(Protocol[_E, _ResT_co]):
    """
    Systems provide functionality that the compiler can target.

    The system will likely involve access to some underlying resource (_ResT_co). There is no strict
    guidance on system object lifetime, however in general, the system object should simply act as a
    facade to the resource object(s). Further, system objects will be managed by runtimes that they
    are added to (see Runtime.add_system), and so the runtime may close the system after a failure
    or if the user has completed their required executions.

    To ensure resource access is appropriately managed, the runtime and binding functions must only
    access resources through the system-provided methods, close and get_resource_proxy.

    Crucially, the process of binding will produce new objects (ExecutableIntermediates and
    Executables) which may have lifetimes detached from their origin runtimes. Hence, resource
    proxies should be implemented in a manner that also prevents unnecessarily holding onto
    resource object references.
    """

    # TODO: Allow multiple systems of the same implementation.
    implements: ClassVar[SupportedSystem]
    exec_mode: ClassVar[ExecutionMode]

    def prepare_resources(self, manifest: Manifest, graph: ManifestGraph) -> None:
        """Prepare system resources before binding any nodes.

        Called once by the runtime before binding begins. Systems can inspect the
        manifest graph and initialize shared resources needed for binding.

        Default implementation is a no-op. Override this method if your system needs
        to initialize resources based on the manifest structure.

        Parameters
        ----------
        manifest : Manifest
            The manifest to be bound.
        graph : ManifestGraph
            The manifest dependency graph for inspection.
        """

    @abc.abstractmethod
    def close(self) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    def bind(self, node: ManifestNode, graph: ManifestGraph, intermediate: _E) -> _E:
        raise NotImplementedError

    @abc.abstractmethod
    def get_resource_proxy(self) -> ResourceProxy[_ResT_co]:
        raise NotImplementedError


_SysT_co = TypeVar("_SysT_co", bound=System, covariant=True)


@runtime_checkable
class SystemInfo(Protocol[_SysT_co]):
    """
    System information object.

    This object should fully describe a system and implement the build_system method to reliably
    create the corresponding system object (_SysT_co) from just data.
    """

    exec_mode: ExecutionMode

    # TODO: Standardize per-system options interface.
    #       See https://qctrl.atlassian.net/browse/SCUP-3201

    @classmethod
    @abc.abstractmethod
    def loads(cls, data: str) -> Self:
        raise NotImplementedError

    @abc.abstractmethod
    def dumps(self) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def build_system(self) -> _SysT_co:
        raise NotImplementedError


@dataclass
class RuntimeOptions:
    """
    Options for runtime behavior.

    Parameters
    ----------
    exec_mode : ExecutionMode
        The execution mode the runtime should operate in. As of implementation, the runtime
        protocols are intended to only support one execution mode per runtime. There is no
        intention to support mixed execution modes.
    close_on_system_error : bool
        Whether to close the runtime and component systems if there are any errors during binding
        and/or adding of new systems.
    close_on_exec_error : bool
        Whether to close the runtime and component systems if an executable bound to this runtime
        raises an error during execution.
    """

    exec_mode: ExecutionMode = ExecutionMode.SYNC

    close_on_system_error: bool = True

    close_on_exec_error: bool = False


class Runtime(Protocol[_E]):
    @property
    def is_closed(self) -> bool:
        raise NotImplementedError

    @property
    def systems(self) -> dict[str, System]:
        raise NotImplementedError

    @classmethod
    @abc.abstractmethod
    def new(cls, options: RuntimeOptions) -> Self:
        raise NotImplementedError

    @abc.abstractmethod
    def close(self) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    def add_system(self, system: System[_E, Any] | SystemInfo) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    def bind(self, manifest: Manifest) -> Executable[_E]:
        raise NotImplementedError

    def run(self, manifest: Manifest) -> ExecutableResult:
        return self.bind(manifest).run()
