# 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__ = (
    "BaseOPXSystem",
    "IQCCExecutionContext",
    "IQCCSystemInfo",
    "OPXExecutionContext",
    "OPXOctaveInfo",
    "OPXSystemInfo",
    "SyncOPX1000System",
    "SyncOPXSystem",
)

import contextlib
import json
import logging
import weakref
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast

import qm
from boulderopalscaleupsdk.runtime.manifest import (
    ConcatenateProgramResultsNode,
    ManifestNode,
    MixerCalibrationNode,
    QuaProgramNode,
    SupportedSystem,
)
from iqcc_cloud_client.qmm_cloud import CloudQuantumMachine, CloudQuantumMachinesManager
from pydantic import BaseModel, Field
from qm.octave import QmOctaveConfig

from boulderopalscaleup.runtime._sync.executable import ExecUnit, NoResourceProxy
from boulderopalscaleup.runtime._sync.quantum_machines import (
    exec_mixer_calibration,
    exec_qua_program_opx,
    exec_qua_program_opx1000,
    exec_standalone_mixer_calibration,
)
from boulderopalscaleup.runtime.abstract import ExecutionMode

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator

    import numpy as np
    from boulderopalscaleupsdk.device.controller.quantum_machines import QuaProgram
    from boulderopalscaleupsdk.runtime.manifest import Manifest
    from pydantic.main import IncEx
    from qm.type_hinting.config_types import ControllerQuaConfig, FullQuaConfig

    from boulderopalscaleup.runtime.abstract import ExecutableIntermediate
    from boulderopalscaleup.runtime.manifest import ManifestGraph


QMM_T = TypeVar("QMM_T", bound=qm.QuantumMachinesManager | CloudQuantumMachinesManager)
QM_T = TypeVar("QM_T", bound=qm.QuantumMachine | CloudQuantumMachine)
CXT_T = TypeVar("CXT_T", bound="IQCCExecutionContext | OPXExecutionContext")


LOGGER = logging.getLogger(__name__)

### OPX+ ###


class OPXExecutionContext(BaseModel):
    """Runtime execution context for OPX+."""


class OPXOctaveInfo(BaseModel):
    """Information for an Octave device.

    Attributes
    ----------
    name : str
        Octave device name.
    host : str
        Hostname or IP address of the Octave.
    port : int
        Port number for Octave connection.
    """

    name: str
    host: str
    port: int


class OPXSystemInfo(BaseModel):
    """System information for OPX controllers.

    Attributes
    ----------
    exec_mode : ExecutionMode
        Execution mode (currently only SYNC supported).
    host : str
        Hostname or IP address of the Quantum Machines Manager.
    port : int
        Port number for QMM connection.
    cluster_name : str or None, optional
        OPX cluster name.
    timeout : int or None, optional
        Connection timeout in seconds.
    octaves : list[OPXOctaveInfo]
        List of Octave devices info.
    log_level : int or None, optional
        Logging level for the Quantum Machines Manager (e.g. ``logging.INFO``).
    octave_calibration_db_path : str or None, optional
        Path to the Octave calibration database directory.
    """

    exec_mode: ExecutionMode = ExecutionMode.SYNC
    host: str
    port: int
    cluster_name: str | None = None
    timeout: int | None = None
    octaves: list[OPXOctaveInfo] = []
    log_level: int | None = None
    octave_calibration_db_path: str | None = None

    @classmethod
    def loads(cls, data: str) -> OPXSystemInfo:
        return cls.model_validate_json(data)

    def dumps(self) -> str:
        return self.model_dump_json()

    def build_system(self) -> SyncOPXSystem:
        octave_config = None
        if self.octaves:
            octave_config = QmOctaveConfig()
            for octave in self.octaves:
                octave_config.add_device_info(
                    name=octave.name,
                    host=octave.host,
                    port=octave.port,
                )

        kwargs: dict = {
            "host": self.host,
            "port": self.port,
            "octave": octave_config,
        }
        if self.cluster_name is not None:
            kwargs["cluster_name"] = self.cluster_name
        if self.timeout is not None:
            kwargs["timeout"] = self.timeout
        if self.log_level is not None:
            kwargs["log_level"] = self.log_level
        if self.octave_calibration_db_path is not None:
            kwargs["octave_calibration_db_path"] = self.octave_calibration_db_path

        qmm = qm.QuantumMachinesManager(**kwargs)

        return SyncOPXSystem(qmm=qmm, context=OPXExecutionContext())


### IQCC / OPX1000 ###


class IQCCExecutionContext(BaseModel):
    """Runtime execution context for Quantum Machines systems.

    Attributes
    ----------
    timeout : float
        Timeout for cloud execution in seconds.
    exclude_result_keys : list[str]
        Result keys to exclude from output (e.g., '__fridge_info').
    """

    timeout: float
    exclude_result_keys: list[str]


class IQCCSystemInfo(BaseModel):
    """System information for IQCC backends.

    Attributes
    ----------
    backend : str
        IQCC backend name (e.g., "meron").
    timeout : float
        Timeout for cloud execution in seconds.
    exclude_result_keys : list[str]
        Result keys to exclude from output.
    """

    exec_mode: ExecutionMode = ExecutionMode.SYNC
    backend: str

    # https://github.com/qctrl/boulder-opal-scale-up/pull/2720#discussion_r2789739789
    timeout: float = 600.0
    exclude_result_keys: list[str] = Field(default_factory=lambda: ["__fridge_info"])

    @classmethod
    def loads(cls, data: str) -> IQCCSystemInfo:
        return cls.model_validate_json(data)

    def dumps(self) -> str:
        return self.model_dump_json()

    def build_system(self) -> SyncOPX1000System:
        qmm = CloudQuantumMachinesManager(backend=self.backend)
        return SyncOPX1000System(
            qmm=qmm,
            context=IQCCExecutionContext.model_validate(
                self.model_dump(include={"timeout", "exclude_result_keys"}),
            ),
        )


class _QMProxy(Generic[QM_T]):
    """Resource proxy for Quantum Machine.

    Uses weak reference to prevent executables from extending system lifetime.
    """

    def __init__(self, sys_ref: weakref.ReferenceType[BaseOPXSystem[QMM_T, QM_T, CXT_T]]):
        self._sys = sys_ref

    @contextlib.contextmanager
    def get(self) -> Iterator[QM_T]:
        """Get managed access to the QM resource.

        Yields
        ------
        QM_T
            The Quantum Machine instance.

        Raises
        ------
        RuntimeError
            If the system has been garbage collected or closed.
        """
        sys = self._sys()
        if sys is None:  # pragma: no cover
            raise RuntimeError("Resource proxy for QM is out-of-bounds.")
        if sys._qm_instance is None:
            raise RuntimeError("OPX system has no QM instance.")

        yield sys._qm_instance


class _QMMProxy(Generic[QMM_T]):
    """Resource proxy for Quantum Machines Manager."""

    def __init__(self, sys_ref: weakref.ReferenceType[BaseOPXSystem[QMM_T, QM_T, CXT_T]]):
        self._sys = sys_ref

    @contextlib.contextmanager
    def get(self) -> Iterator[QMM_T]:
        """Get managed access to the QMM resource.

        Yields
        ------
        QMM_T
            The Quantum Machines Manager instance.

        Raises
        ------
        RuntimeError
            If the system has been garbage collected or closed.
        """
        sys = self._sys()
        if sys is None:  # pragma: no cover
            raise RuntimeError("Resource proxy for QMM is out-of-bounds.")
        if sys._qmm is None:
            raise RuntimeError("OPX system has been closed.")

        yield sys._qmm


class BaseOPXSystem(ABC, Generic[QMM_T, QM_T, CXT_T]):
    """Base class for OPX systems.

    Attributes
    ----------
    implements : SupportedSystem
        Identifies this as a Quantum Machines system.
    exec_mode : ExecutionMode
        Synchronous execution mode.
    """

    implements: ClassVar[SupportedSystem] = SupportedSystem.QUANTUM_MACHINES
    exec_mode: ClassVar[ExecutionMode] = ExecutionMode.SYNC

    def __init__(
        self,
        qmm: QMM_T,
        context: CXT_T,
    ):
        """Initialize base OPX system.

        Parameters
        ----------
        qmm : QMM_T
            Quantum Machines Manager instance.
        context : CXT_T
            Runtime context data.
        """
        self._qmm: QMM_T | None = qmm
        self._qm_instance: QM_T | None = None
        self._context = context

    def close(self) -> None:
        """Close the system and release resources."""
        if self._qm_instance is not None:
            try:
                self._qm_instance.close()
            except Exception as e:  # noqa: BLE001
                LOGGER.warning("Failed to close QM instance: %s", e)
            self._qm_instance = None
        self._qmm = None

    def get_resource_proxy(self) -> _QMProxy[QM_T]:
        """Get resource proxy for Quantum Machine."""
        return _QMProxy(weakref.ref(self))

    def get_qmm_proxy(self) -> _QMMProxy[QMM_T]:
        """Get resource proxy for Quantum Machines Manager."""
        return _QMMProxy(weakref.ref(self))

    def prepare_resources(
        self,
        manifest: Manifest,
        graph: ManifestGraph,  # noqa: ARG002
    ) -> None:
        """Prepare QM instance for execution.

        Parameters
        ----------
        manifest : Manifest
            The complete manifest to be bound.
        graph : ManifestGraph
            The manifest dependency graph.
        """
        if self._qmm is None:
            raise RuntimeError("Cannot prepare resources: OPX system has been closed.")

        qua_program_nodes = [node for node in manifest.nodes if isinstance(node, QuaProgramNode)]

        if not qua_program_nodes:
            # No QUA programs in manifest, nothing to prepare
            return

        # Open QM with the first program's config
        # TODO: Handle multiple QuaProgramNodes with different configs.
        # For now we assume all programs use the same config and just use the first one we find.
        first_program = qua_program_nodes[0]
        config_json = first_program.data.config.root.model_dump_json(
            exclude=self._get_config_exclude(),
            exclude_none=True,
        )
        qua_config = json.loads(config_json)
        self._qm_instance = self._open_qm(qua_config)

        if len(qua_program_nodes) > 1:
            LOGGER.warning(
                "Manifest contains %d QuaProgramNodes. Using config from first node '%s'. "
                "Multiple configs are not validated for compatibility.",
                len(qua_program_nodes),
                first_program.id,
            )

    @abstractmethod
    def _get_config_exclude(self) -> IncEx:
        """Get config exclusion pattern for this OPX variant.

        Returns
        -------
        IncEx
            Pydantic exclusion pattern for config serialization.
        """
        raise NotImplementedError

    @abstractmethod
    def _get_exec_func(
        self,
    ) -> Callable[[QuaProgram, QM_T, CXT_T], dict[str, np.ndarray]]:
        """Get execution function for this OPX variant.

        Returns
        -------
        Callable
            Execution function for running QUA programs.
        """
        raise NotImplementedError

    @abstractmethod
    def _open_qm(self, config: dict) -> QM_T:
        """Open quantum machine with system-specific configuration.

        Parameters
        ----------
        config : dict
            QUA configuration dictionary.

        Returns
        -------
        QM_T
            Opened quantum machine instance.
        """
        raise NotImplementedError

    def bind(
        self,
        node: ManifestNode,
        graph: ManifestGraph,  # noqa: ARG002
        intermediate: ExecutableIntermediate[ExecUnit],
    ) -> ExecutableIntermediate[ExecUnit]:
        """Bind a manifest node to an execution unit.

        Parameters
        ----------
        node : ManifestNode
            The manifest node to bind.
        _ : ManifestGraph
            The complete manifest graph for context.
        intermediate : ExecutableIntermediate[ExecUnit]
            The intermediate executable being constructed.

        Returns
        -------
        ExecutableIntermediate[ExecUnit]
            Updated intermediate with bound unit.

        Raises
        ------
        TypeError
            If node type is not supported by this system.
        """
        match node:
            case QuaProgramNode():
                self._bind_qua_program(node, intermediate)
            case MixerCalibrationNode():
                self._bind_mixer_calibration(node, intermediate)
            case ConcatenateProgramResultsNode():
                # TODO: Handle concatenation of multiple Qua programs when needed.
                ...
            case _:  # pragma: no cover
                raise TypeError(
                    f"{type(self).__name__} does not support node of type {type(node).__name__}",
                )
        return intermediate

    @abstractmethod
    def _bind_mixer_calibration(
        self,
        node: MixerCalibrationNode,
        intermediate: ExecutableIntermediate[ExecUnit],
    ) -> None:
        raise NotImplementedError

    def _bind_qua_program(
        self,
        node: QuaProgramNode,
        intermediate: ExecutableIntermediate[ExecUnit],
    ) -> None:
        """Bind a QuaProgramNode to an execution unit.

        Parameters
        ----------
        node : QuaProgramNode
            The Qua program node to bind.
        intermediate : ExecutableIntermediate[ExecUnit]
            The intermediate executable being constructed.
        """
        if self._qmm is None:
            raise RuntimeError("Cannot bind program: OPX system has been closed.")

        if self._qm_instance is None:
            raise RuntimeError("Quantum Machine instance not prepared.")

        unit = ExecUnit(
            id=node.id,
            data=node.data,
            resource_proxy=self.get_resource_proxy(),
            context=self._context,
            run_func=self._get_exec_func(),
            tracks_duration=True,
        )
        intermediate.set_unit(node.id, unit)


class SyncOPXSystem(
    BaseOPXSystem[qm.QuantumMachinesManager, qm.QuantumMachine, OPXExecutionContext],
):
    """Synchronous execution system for OPX+ controllers."""

    def _get_config_exclude(self) -> IncEx:
        """Get config exclusion pattern for OPX."""
        return {"qm_version"}

    def _get_exec_func(self) -> Callable:
        """Get execution function for OPX."""
        return exec_qua_program_opx

    def _open_qm(  # pyright: ignore[reportIncompatibleMethodOverride]
        self,
        config: FullQuaConfig | ControllerQuaConfig,  # type: ignore[override]
    ) -> qm.QuantumMachine:
        """Open quantum machine for OPX.

        Parameters
        ----------
        config : dict
            QUA configuration dictionary.

        Returns
        -------
        qm.QuantumMachine
            Opened quantum machine instance.
        """
        if self._qmm is None:
            raise RuntimeError("Cannot open QM: OPX system has been closed.")
        return self._qmm.open_qm(config)  # type: ignore[return-value]

    def _bind_mixer_calibration(
        self,
        node: MixerCalibrationNode,
        intermediate: ExecutableIntermediate[ExecUnit],
    ) -> None:
        """Bind a MixerCalibrationNode to an execution unit."""
        if self._qmm is None:
            raise RuntimeError("Cannot bind calibration: OPX system has been closed.")

        if node.data.config is not None:
            intermediate.set_unit(
                node.id,
                ExecUnit(
                    id=node.id,
                    data=node.data,
                    resource_proxy=self.get_qmm_proxy(),
                    context=self._context,
                    run_func=exec_standalone_mixer_calibration,
                    tracks_duration=True,
                ),
            )
        else:
            if self._qm_instance is None:
                raise RuntimeError("Quantum Machine instance not prepared.")

            intermediate.set_unit(
                node.id,
                ExecUnit(
                    id=node.id,
                    data=node.data.elements,
                    resource_proxy=self.get_resource_proxy(),
                    context=self._context,
                    run_func=exec_mixer_calibration,
                    tracks_duration=True,
                ),
            )


class SyncOPX1000System(
    BaseOPXSystem[CloudQuantumMachinesManager, CloudQuantumMachine, IQCCExecutionContext],
):
    """Synchronous execution system for OPX1000 controllers.

    Attributes
    ----------
    implements : SupportedSystem
        Identifies this as a Quantum Machines system (inherited).
    exec_mode : ExecutionMode
        Synchronous execution mode (inherited).
    """

    def _get_config_exclude(self) -> IncEx:
        """Get config exclusion pattern for OPX1000.

        Returns
        -------
        IncEx
            Excludes qm_version and feedback filter fields for QOP 3.5 compatibility.
        """
        return cast(
            "IncEx",
            {
                "qm_version": True,
                "controllers": {
                    "__all__": {
                        "fems": {
                            "__all__": {
                                "analog_outputs": {
                                    "__all__": {
                                        "filter": {
                                            "feedback": True,
                                        },
                                    },
                                },
                            },
                        },
                    },
                },
            },
        )

    def _get_exec_func(self) -> Callable:
        """Get execution function for OPX1000.

        Returns
        -------
        Callable
            OPX1000 execution function.
        """
        return exec_qua_program_opx1000

    def _open_qm(self, config: dict) -> CloudQuantumMachine:
        """Open quantum machine for OPX1000.

        Parameters
        ----------
        config : dict
            QUA configuration dictionary.

        Returns
        -------
        CloudQuantumMachine
            Opened cloud quantum machine instance.
        """
        if self._qmm is None:
            raise RuntimeError("Cannot open QM: OPX system has been closed.")
        return self._qmm.open_qm(config=config, close_other_machines=True)

    def _bind_mixer_calibration(
        self,
        node: MixerCalibrationNode,
        intermediate: ExecutableIntermediate[ExecUnit],
    ) -> None:
        LOGGER.warning("Mixer calibration not required. Skipping for node %s on OPX1000.", node.id)
        intermediate.set_unit(
            node.id,
            ExecUnit(
                id=node.id,
                data=node.data.elements,
                resource_proxy=NoResourceProxy(),
                context=self._context,
                run_func=lambda _data, _resource, _ctx: {},
            ),
        )
