# 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

import json
import logging
from typing import TYPE_CHECKING

import numpy as np
import qm
from qm.exceptions import AnotherJobIsRunning, CantCalibrateElementError

from boulderopalscaleup.runtime.abstract import CalibrationError
from boulderopalscaleup.utils import sync_retry_with

if TYPE_CHECKING:
    from boulderopalscaleupsdk.device.controller.quantum_machines import QuaProgram
    from boulderopalscaleupsdk.runtime.manifest import MixerCalibrationData
    from iqcc_cloud_client.qmm_cloud import CloudJob, CloudQuantumMachine
    from qm.jobs.running_qm_job import RunningQmJob
    from qm.octave.octave_mixer_calibration import (
        MixerCalibrationResults,
    )

    from boulderopalscaleup.runtime.systems.quantum_machines import (
        IQCCExecutionContext,
        OPXExecutionContext,
    )

LOGGER = logging.getLogger(__name__)


def exec_qua_program_opx(
    qua_program: QuaProgram,
    qm_instance: qm.QuantumMachine,
    ctx: OPXExecutionContext,  # noqa: ARG001
) -> dict[str, list[float]]:
    """Execute a Qua program on OPX and return formatted results.

    Parameters
    ----------
    qua_program : QuaProgram
        The Qua program to execute, containing both program and config.
    qm_instance : qm.QuantumMachine
        The Quantum Machine instance for executing programs.
    ctx : OPXExecutionContext
        Execution context for OPX.

    Returns
    -------
    dict[str, list[float]]
        Formatted results with stream names as keys and lists of floats as values.
    """
    job = qm_instance.execute(qm.Program(program=qua_program.program))

    return _fetch_and_format_results(job)


def exec_qua_program_opx1000(
    qua_program: QuaProgram,
    qm_instance: CloudQuantumMachine,
    ctx: IQCCExecutionContext,
) -> dict[str, list[float]]:
    """Execute a Qua program on OPX1000 and return formatted results.

    Parameters
    ----------
    qua_program : QuaProgram
        The Qua program to execute, containing both program and config.
    qm_instance : CloudQuantumMachine
        The Quantum Machine instance for executing programs.
    ctx : IQCCExecutionContext
        Execution context containing runtime context.
    """

    @sync_retry_with(max_retries=10)
    def _execute() -> CloudJob:
        return qm_instance.execute(
            qm.Program(program=qua_program.program),
            options={"timeout": ctx.timeout},
        )

    job = _execute()

    return _fetch_and_format_results(job, exclude_keys=ctx.exclude_result_keys)


def _fetch_and_format_results(
    job: RunningQmJob | CloudJob,
    exclude_keys: list[str] | None = None,
) -> dict[str, list[float]]:
    """Fetch results from a QM job and format them.

    Parameters
    ----------
    job: RunningQmJob | CloudJob
        The running Quantum Machines job.
    exclude_keys : list[str] | None
        Optional list of result keys to exclude from the output.

    Returns
    -------
    dict[str, list[float]]
        Results with stream names as keys and lists of floats as values.
    """

    def _format(data):
        """Format structured arrays by extracting the 'value' field."""
        if (
            type(data) is np.ndarray
            and data.ndim > 0
            and len(data) > 0
            and type(data[0]) is np.void
            and len(data.dtype.names or []) == 1
        ):
            data = data["value"]
        return data

    if exclude_keys is None:
        exclude_keys = []

    result_handles = job.result_handles
    result_handles.wait_for_all_values()

    results: dict[str, list[float]] = {}
    for stream_name in list(result_handles.keys()):
        if stream_name in exclude_keys:
            continue
        if hasattr(result_handles, stream_name):
            fetched_data = result_handles.get(stream_name).fetch_all()  # type: ignore[union-attr]
            if fetched_data is not None:
                results[stream_name] = np.asarray(_format(fetched_data)).astype(float).tolist()

    return results


def exec_standalone_mixer_calibration(
    data: MixerCalibrationData,
    qmm: qm.QuantumMachinesManager,
    ctx: OPXExecutionContext,  # noqa: ARG001
) -> CalibrationError | dict[str, MixerCalibrationResults]:
    """Execute mixer calibration using a dedicated QM opened from the request config.

    Backwards compatible logic for the existing RunQuantumMachinesMixerCalibrationRequest.

    Parameters
    ----------
    data : MixerCalibrationData
        Calibration data including element names and the QUA config JSON.
    qmm : qm.QuantumMachinesManager
        The Quantum Machines Manager used to open a dedicated QM instance.
    ctx : OPXExecutionContext
        Execution context (currently unused, reserved for future use).
    """
    if data.config is None:
        raise RuntimeError("Cannot run standalone mixer calibration: config is not set.")

    # TODO: manage multiple QM instance lifecycle in OPXSystem
    # Can be consolidated with exec_mixer_calibration as a result
    # https://qctrl.atlassian.net/browse/SCUP-3677
    qm_instance = qmm.open_qm(json.loads(data.config))

    results = {}
    for element in data.elements:
        try:
            result = qm_instance.calibrate_element(element)
            results[element] = result
            LOGGER.info("Successfully calibrated element: %s", element)
        except CantCalibrateElementError as error:
            return CalibrationError(
                message=f"Failed to calibrate element {element}: {error}",
            )
        except AnotherJobIsRunning:
            return CalibrationError(
                message=f"Failed to calibrate element {element}: "
                f"another controller job is running.",
            )
    return results


def exec_mixer_calibration(
    elements: list[str],
    qm_instance: qm.QuantumMachine,
    ctx: OPXExecutionContext,  # noqa: ARG001
) -> CalibrationError | dict[str, MixerCalibrationResults]:
    """Execute mixer calibration on OPX.

    Parameters
    ----------
    elements : list[str]
        List of element names to calibrate.
    qm_instance : qm.QuantumMachine
        The Quantum Machine instance for executing calibration.
    ctx : OPXExecutionContext
        Execution context (currently unused for calibration, reserved for future use).
    """
    results = {}
    for element in elements:
        try:
            result = qm_instance.calibrate_element(element)
            results[element] = result
            LOGGER.info("Successfully calibrated element: %s", element)
        except CantCalibrateElementError as error:
            return CalibrationError(
                message=f"Failed to calibrate element {element}: {error}",
            )
        except AnotherJobIsRunning as error:
            return CalibrationError(
                message=f"Failed to calibrate element {element}: {error}",
            )

    return results
