# 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

from boulderopalscaleup.runtime.abstract import ExecutionMode, ResourceProxy

__all__ = (
    "BaseQbloxSystem",
    "ConcatenateData",
    "PreparedProgramUnitData",
    "QbloxClusterInfo",
    "QbloxExecutionContext",
    "QbloxRoutines",
    "QbloxSystemInfo",
)

import abc
import enum
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar

import qblox_instruments
from boulderopalscaleupsdk.device.controller.qblox import (
    DEFAULT_DUMMY_CLUSTER,
    ModuleAddr,
    ModuleAddrType,
    ModuleType,
    OutputSequencerAcquisitions,
    PreparedProgram,
    SequencerAddr,
)
from boulderopalscaleupsdk.runtime.manifest import (
    ConcatenateProgramResultsNode,
    Manifest,
    ManifestNode,
    MixerCalibrationNode,
    PreparedProgramNode,
    ReshapeProgramResultsNode,
    SupportedSystem,
)
from pydantic import BaseModel, Field
from qcodes.instrument import find_or_create_instrument

if TYPE_CHECKING:
    from boulderopalscaleup.runtime.abstract import (
        ExecutableIntermediate,
        ResultProxy,
    )
    from boulderopalscaleup.runtime.manifest import ManifestGraph

RawT = dict[SequencerAddr, OutputSequencerAcquisitions]


class QbloxExecutionContext(BaseModel):
    """Runtime execution context for QBLOX systems.

    Attributes
    ----------
    timeout : float
        Maximum execution time in seconds before timeout.
    timeout_poll_res : float
        Polling resolution in seconds for checking execution status.
    """

    timeout: float
    timeout_poll_res: float


class QbloxClusterInfo(BaseModel):
    """Connection information for a single Qblox cluster.

    Attributes
    ----------
    name : str
        Instrument name used to identify the cluster.
    host : str
        Hostname or IP address of the cluster.
    """

    name: str
    host: str


class QbloxDummySystem(BaseModel):
    """Dummy cluster used for compilation testing."""

    exec_mode: ExecutionMode = ExecutionMode.SYNC
    name: str = DEFAULT_DUMMY_CLUSTER
    modules: dict[ModuleAddrType, ModuleType] = Field(
        default={
            ModuleAddr(cluster=DEFAULT_DUMMY_CLUSTER, slot=slot): ModuleType.QRM_RF
            for slot in range(20)
        },
    )
    timeout: float = 30.0
    timeout_poll_res: float = 1.0

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

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

    @staticmethod
    def _cluster_type(mod_type: ModuleType) -> qblox_instruments.ClusterType:
        match mod_type:
            case ModuleType.QCM:
                return qblox_instruments.ClusterType.CLUSTER_QCM
            case ModuleType.QCM_RF:
                return qblox_instruments.ClusterType.CLUSTER_QCM_RF
            case ModuleType.QRM:
                return qblox_instruments.ClusterType.CLUSTER_QRM
            case ModuleType.QRM_RF:
                return qblox_instruments.ClusterType.CLUSTER_QRM_RF
            case _:  # pragma: no cover
                raise ValueError(f"Cannot create dummy cluster with module of type {mod_type!s}.")

    def build_system(self) -> BaseQbloxSystem:
        try:
            cls = _REGISTERED_SYSTEM_IMPL[self.exec_mode]
        except KeyError:  # pragma: no cover
            raise ValueError(
                f"QBLOX system does not support execution mode '{self.exec_mode}'",
            ) from None
        stack = {
            self.name: find_or_create_instrument(
                qblox_instruments.Cluster,
                recreate=True,
                name=self.name,
                identifier=None,
                dummy_cfg={
                    mod.slot: self._cluster_type(mod_type) for mod, mod_type in self.modules.items()
                },
            ),
        }
        context = QbloxExecutionContext.model_validate(
            self.model_dump(include={"timeout", "timeout_poll_res"}),
        )
        return cls(stack, context)


class QbloxSystemInfo(BaseModel):
    """System information for a Qblox hardware backend.

    Describes one or more Qblox clusters to connect to and the execution
    parameters to use when running programs. Pass to ``Runtime.add_system()``
    to register the Qblox backend with a runtime.

    Attributes
    ----------
    exec_mode : ExecutionMode
        Internal execution mode. Defaults to `ExecutionMode.SYNC`.
    clusters : list[QbloxClusterInfo]
        List of clusters to connect to.
    timeout : float
        Maximum execution time in seconds before a timeout error is raised.
        Defaults to ``30.0``.
    timeout_poll_res : float
        Polling interval in seconds for checking execution status.
        Defaults to ``1.0``.
    """

    exec_mode: ExecutionMode = ExecutionMode.SYNC
    clusters: list[QbloxClusterInfo]
    timeout: float = 30.0
    timeout_poll_res: float = 1.0

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

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

    def build_system(self) -> BaseQbloxSystem:
        try:
            cls = _REGISTERED_SYSTEM_IMPL[self.exec_mode]
        except KeyError:  # pragma: no cover
            raise ValueError(
                f"QBLOX system does not support execution mode '{self.exec_mode}'",
            ) from None
        stack = {
            cc.name: find_or_create_instrument(qblox_instruments.Cluster, cc.name, cc.host)
            for cc in self.clusters
        }
        context = QbloxExecutionContext.model_validate(
            self.model_dump(include={"timeout", "timeout_poll_res"}),
        )
        return cls(stack, context)


class QbloxRoutines(enum.Enum):
    RUN_PROGRAM_RAW = enum.auto()
    RUN_PROGRAM_PROCESSED = enum.auto()
    CONCAT = enum.auto()
    RESHAPE = enum.auto()
    MIXER_CALIBRATION = enum.auto()


@dataclass
class ConcatenateData:
    """Data required for concatenating program results.

    Implements ResultsInjectable protocol for type-safe results injection.
    """

    dep_proxies: dict[str, ResultProxy[RawT]]
    dep_programs: dict[str, PreparedProgram]


@dataclass
class ReshapeData:
    """Data required for reshaping program results into a multi-dimensional array."""

    dep_proxies: dict[str, ResultProxy[RawT]]
    dep_programs: dict[str, PreparedProgram]
    result_shape: tuple[int, ...]


@dataclass
class PreparedProgramUnitData:
    """Data required for program execution with optional calibration.

    Attributes
    ----------
    prepared_program : PreparedProgram
        The prepared program to execute.
    calibrate_elements : list[str] | None
        Optional list of element names to calibrate.
    """

    prepared_program: PreparedProgram
    calibrate_elements: list[str] | None = None


_UnitT = TypeVar("_UnitT")
_REGISTERED_SYSTEM_IMPL: dict[ExecutionMode, type[BaseQbloxSystem]] = {}


class BaseQbloxSystem(abc.ABC, Generic[_UnitT]):
    implements: ClassVar[SupportedSystem] = SupportedSystem.QBLOX
    exec_mode: ClassVar[ExecutionMode]

    _stack: dict[str, qblox_instruments.Cluster] | None
    _context: QbloxExecutionContext

    def __init__(
        self,
        stack: dict[str, qblox_instruments.Cluster],
        context: QbloxExecutionContext,
    ):
        self._stack = stack
        self._context = context

    def __init_subclass__(cls, execution_mode: ExecutionMode, *args, **kwargs):
        super().__init_subclass__(*args, **kwargs)
        cls.exec_mode = execution_mode
        _REGISTERED_SYSTEM_IMPL[execution_mode] = cls

    def close(self) -> None:
        if self._stack is None:
            return
        for cluster in self._stack.values():
            cluster.close()
        self._stack = None

    def prepare_resources(
        self,
        manifest: Manifest,
        graph: ManifestGraph,
    ) -> None:
        pass

    @abc.abstractmethod
    def make_unit(self, node_id: str, data: Any, routine: QbloxRoutines) -> _UnitT:
        raise NotImplementedError

    @abc.abstractmethod
    def get_resource_proxy(self) -> ResourceProxy[dict[str, qblox_instruments.Cluster]]:
        raise NotImplementedError

    def bind(
        self,
        node: ManifestNode,
        graph: ManifestGraph,
        intermediate: ExecutableIntermediate[_UnitT],
    ) -> ExecutableIntermediate[_UnitT]:
        match node:
            case PreparedProgramNode():
                self._bind_prepared_program(node, graph, intermediate)
            case ConcatenateProgramResultsNode():
                self._bind_concatenate_results(node, graph, intermediate)
            case ReshapeProgramResultsNode():
                self._bind_reshape_results(node, graph, intermediate)
            case MixerCalibrationNode():
                self._bind_mixer_calibration(node, graph, intermediate)
            case _:  # pragma: no cover
                raise TypeError(
                    f"QBLOX runtime does not support node of type {type(node).__name__} ",
                )
        return intermediate

    def _does_sinks_to_post_processing(
        self,
        node: PreparedProgramNode,
        graph: ManifestGraph,
    ) -> bool:
        """Check if the node eventually sinks to a post-processing node."""
        all_paths = graph.find_paths(node.id)
        for path in all_paths:
            for nn in path:
                next_node = graph.nodes[nn]
                match next_node:
                    case PreparedProgramNode():
                        continue  # Allow program-program dependency
                    case ConcatenateProgramResultsNode() | ReshapeProgramResultsNode():
                        # TODO: This doesn't handle the case where the program is sunk to
                        #       DIFFERENT paths!
                        #       We actually have to place a post-processing function in one
                        #       path, while we allow the raw results to flow to others.
                        return True
                    case _:
                        break  # Program does not sink into Concat/Reshape...
        return False

    def _bind_prepared_program(
        self,
        node: PreparedProgramNode,
        graph: ManifestGraph,
        intermediate: ExecutableIntermediate[_UnitT],
    ) -> None:
        """Bind a PreparedProgramNode to an execution unit."""
        calibrate_elements: list[str] | None = None
        for dep_id in node.depends_on:
            dep_node = graph.nodes.get(dep_id)
            if isinstance(dep_node, MixerCalibrationNode):
                calibrate_elements = dep_node.data.elements
                break

        unit_data = PreparedProgramUnitData(
            prepared_program=node.data,
            calibrate_elements=calibrate_elements,
        )

        if self._does_sinks_to_post_processing(node, graph):
            unit = self.make_unit(node.id, unit_data, QbloxRoutines.RUN_PROGRAM_RAW)
        else:
            unit = self.make_unit(node.id, unit_data, QbloxRoutines.RUN_PROGRAM_PROCESSED)
        intermediate.set_unit(node.id, unit)

    def _bind_concatenate_results(
        self,
        node: ConcatenateProgramResultsNode,
        graph: ManifestGraph,
        intermediate: ExecutableIntermediate[_UnitT],
    ) -> None:
        """Bind a ConcatenateProgramResultsNode to an execution unit."""
        dep_programs: dict[str, PreparedProgram] = {}
        dep_proxies: dict[str, ResultProxy[RawT]] = {}
        for dep_id in node.depends_on:
            dep_node = graph.nodes[dep_id]

            # TODO: What if dep_node is not a manifest with `data`?
            # ConcatenateProgramResultsNode for example
            if not isinstance(dep_node, PreparedProgramNode):
                raise TypeError(
                    f"ConcatenateProgramResultsNode '{node.id}' depends on "
                    f"'{dep_id}' which is not a PreparedProgramNode.",
                )

            dep_programs[dep_id] = dep_node.data
            dep_proxies[dep_id] = intermediate.get_result_proxy(dep_id)

        unit = self.make_unit(
            node.id,
            ConcatenateData(dep_proxies, dep_programs),
            QbloxRoutines.CONCAT,
        )
        intermediate.set_unit(node.id, unit)

    def _bind_reshape_results(
        self,
        node: ReshapeProgramResultsNode,
        graph: ManifestGraph,
        intermediate: ExecutableIntermediate[_UnitT],
    ) -> None:
        """Bind a ReshapeProgramResultsNode to an execution unit."""
        dep_programs: dict[str, PreparedProgram] = {}
        dep_proxies: dict[str, ResultProxy[RawT]] = {}
        for dep_id in node.depends_on:
            dep_node = graph.nodes[dep_id]
            if not isinstance(dep_node, PreparedProgramNode):
                raise TypeError(
                    f"ReshapeProgramResultsNode '{node.id}' depends on "
                    f"'{dep_id}' which is not a PreparedProgramNode.",
                )
            dep_programs[dep_id] = dep_node.data
            dep_proxies[dep_id] = intermediate.get_result_proxy(dep_id)

        unit = self.make_unit(
            node.id,
            ReshapeData(dep_proxies, dep_programs, node.result_shape),
            QbloxRoutines.RESHAPE,
        )
        intermediate.set_unit(node.id, unit)

    def _bind_mixer_calibration(
        self,
        node: MixerCalibrationNode,
        graph: ManifestGraph,  # noqa: ARG002
        intermediate: ExecutableIntermediate[_UnitT],
    ) -> None:
        """Acknowledge MixerCalibrationNode and create a no-op execution unit.

        MixerCalibrationNode acts as a marker in the manifest dependency graph.
        The actual mixer calibration is performed by `arm_sequencers` when
        PreparedProgramNode extracts this node's data.
        """
        unit = self.make_unit(
            node.id,
            node.data,
            QbloxRoutines.MIXER_CALIBRATION,
        )
        intermediate.set_unit(node.id, unit)
