# 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__ = ("ExecUnit", "NoResourceProxy", "SyncExecIntermediate")

import contextlib
import graphlib
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Generic, Self, TypeVar

from boulderopalscaleup.runtime.abstract import (
    ExecutableIntermediate,
    ExecutionMetadata,
    ResourceProxy,
    ResultProxy,
)
from boulderopalscaleup.runtime.manifest import ManifestGraph
from boulderopalscaleup.utils import trace

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator

    from boulderopalscaleupsdk.runtime.manifest import Manifest

_DataT = TypeVar("_DataT")
_ResultT = TypeVar("_ResultT")
_ResT = TypeVar("_ResT")
_CtxT = TypeVar("_CtxT")


class _NotSet: ...


class _ResultProxy(Generic[_ResultT]):
    def __init__(self, node_id: str) -> None:
        self._id = node_id
        self._value: _ResultT | _NotSet = _NotSet()

    def set(self, value: _ResultT) -> None:
        self._value = value

    def get(self) -> _ResultT:
        if isinstance(self._value, _NotSet):  # pragma: no cover
            raise RuntimeError(  # noqa: TRY004
                f"Execution for '{self._id}' not completed but attempted to access result.",
            )
        return self._value


@dataclass
class ExecUnit(Generic[_DataT, _ResT, _CtxT, _ResultT]):
    """
    A unit of execution at runtime.

    Parameters
    ----------
    id
        The id of the unit.
    data
        The program data required for execution.
    resource_proxy
        A proxy that provides access to the resource for execution.
    context
        Execution context containing runtime configuration and state.
    run_func
        The function that executes the unit.
    """

    id: str
    data: _DataT
    resource_proxy: ResourceProxy[_ResT]
    context: _CtxT
    run_func: Callable[[_DataT, _ResT, _CtxT], _ResultT]
    # TODO: Consider adding this type of internal config to `context` as well
    tracks_duration: bool = False


@dataclass
class _Result:
    _output: Any
    _intermediates: dict[str, Any]
    _metadata: ExecutionMetadata

    @property
    def output(self) -> Any:
        return self._output

    @property
    def metadata(self) -> ExecutionMetadata:
        return self._metadata

    def get_intermediate(self, node_id: str) -> Any:
        return self._intermediates[node_id]


@dataclass
class SyncExecIntermediate(ExecutableIntermediate[ExecUnit]):
    _units: dict[str, ExecUnit]  # node ID -> ExecUnit
    _graph: ManifestGraph
    _return_node: str

    _result_proxies: dict[str, _ResultProxy] = field(
        init=False,
        default_factory=dict,
    )  # node ID -> proxy

    def get_result_proxy(self, node_id: str) -> ResultProxy:
        if node_id in self._result_proxies:
            return self._result_proxies[node_id]

        proxy = _ResultProxy(node_id)  # type: ignore[var-annotated]
        self._result_proxies[node_id] = proxy
        return proxy

    def set_unit(self, node_id, unit: ExecUnit) -> None:
        if node_id in self._units:
            raise ValueError(f"Node '{node_id}' already has a unit.")
        self._units[node_id] = unit

    def run(self) -> _Result:
        """Execute all units in topological order and return results."""

        # TODO: Safer typing for results?

        intermediates = {}
        sorter = graphlib.TopologicalSorter(self._graph.predecessors)
        ret = _NotSet()
        total_duration_ms: float = 0.0

        for node_id in sorter.static_order():
            unit = self._units.get(node_id)
            if unit is None:  # pragma: no cover
                raise RuntimeError(f"No unit bound for node '{node_id}'.")

            with unit.resource_proxy.get() as resource:
                if unit.tracks_duration:
                    result, duration_ms = trace(unit.run_func, unit.data, resource, unit.context)
                    total_duration_ms += duration_ms
                else:
                    result = unit.run_func(unit.data, resource, unit.context)

            if unit.id in self._result_proxies:
                self._result_proxies[unit.id].set(result)

            if unit.id == self._return_node:
                ret = result

            intermediates[unit.id] = result

        if isinstance(ret, _NotSet):  # pragma: no cover
            raise RuntimeError(f"Required end node '{self._return_node}' never ran")  # noqa: TRY004

        return _Result(
            ret,
            intermediates,
            ExecutionMetadata(duration=total_duration_ms),
        )

    @classmethod
    def init_from_manifest(cls, manifest: Manifest) -> Self:
        return cls({}, ManifestGraph.from_manifest(manifest), manifest.end_node)


class NoResourceProxy(ResourceProxy[Any]):
    """A resource proxy that provides no resources. Used for nodes that don't require hardware."""

    @contextlib.contextmanager
    def get(self) -> Iterator[None]:
        yield None
