"""
Module containing the ArbokDriver class for managing and running sequences
on a physical OPX instrument.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import copy

import xarray as xr
from qm import SimulationConfig, generate_qua_script
from qm.quantum_machines_manager import QuantumMachinesManager
from qcodes.instrument import Instrument
from qcodes.dataset import load_or_create_experiment
from sqlalchemy.orm import Session

from .device import Device
from .measurement import Measurement
from .sqlalchemy_classes import SqlRun

if TYPE_CHECKING:
    from .experiment import Experiment
    from .sqlalchemy_classes import SqlRun
    from qm.jobs.running_qm_job import RunningQmJob
    from qm.quantum_machine import QuantumMachine
    from qm.quantum_machines_manager import QuantumMachinesManager
    from qm import StreamsManager

class ArbokDriver(Instrument):
    """
    Class containing all functionality to manage and run modular sequences on a 
    physical OPX instrument
    """
    opx: QuantumMachine
    qm_job: RunningQmJob
    qmm: QuantumMachinesManager
    result_handles: StreamsManager

    def __init__(
            self,
            name: str,
            device: Device,
            **kwargs
            ) -> None:
        """
        Constructor class for `Program` class
        
        Args:
            name (str): Name of the instrument
            device (Device): Device class describing phyical device
            **kwargs: Arbitrary keyword arguments for qcodes Instrument class
        """
        super().__init__(name, **kwargs)
        if not isinstance(device, Device):
            raise TypeError(f"device must be type Device, is {type(Device)}")
        self.device: Device = device
        self.is_mock: bool = False
        self._measurements: list[Measurement] = []
        self.add_parameter('iteration', get_cmd = None, set_cmd =None)

        self.database_engine = None
        self.minio_filesystem = None
        self.station = None

    @property
    def measurements(self) -> list[Measurement]:
        """Measurements to be run within program uploaded to the OPX"""
        return self._measurements

    def reset_measurements(self) -> None:
        """
        Resets all measurements in the program
        TODO: delete instances of those measurements
        """
        for measurement in self._measurements:
            print(f"Deleting measurement: {measurement.short_name}")
            del measurement
        self._measurements = []
        self.submodules = {}

    def connect_opx(
            self,
            host_ip: str,
            qm_config: dict = None,
            reconnect: bool = False,
            **kwargs) -> None:
        """
        Creates QuantumMachinesManager and opens a quantum machine on it with
        the given IP address
        
        Args:
            host_ip (str): Ip address of the OPX
            qm_config (dict): QM/OPX config dictionary to be used. Defaults to
                None, in which case the config from the device is used. If given
                overwrites the device config.
            reconnect (bool): Whether to reconnect to the OPX if already
                connected. Keeps QMM alive if true. Defaults to False.
            **kwargs: Arbitrary keyword arguments for QMM instanciation
        """
        if not reconnect:
            self.qmm = QuantumMachinesManager(
                host = host_ip, **kwargs)
        if qm_config is not None:
            if not isinstance(qm_config, dict):
                raise ValueError(
                    "qm_config must be a dictionary, not a string")
            self.device.config = copy.deepcopy(qm_config)
        self.opx = self.qmm.open_qm(self.device.config)

    def reconnect_opx(
            self, host_ip: str, qm_config: dict = None) -> None:
        """
        Reconnects to the OPX with the given IP address and closes the previous
        connection
        
        Args:
            host_ip (str): Ip address of the OPX
            qm_config (dict): QM/OPX config dictionary to be used. Defaults to
                None, in which case the config from the device is used. If given
                overwrites the device config. 
        """
        if self.opx is not None:
            print('Closing previous connection')
            self.opx.close()
            self.qmm.close_all_quantum_machines()
        self.connect_opx(host_ip, qm_config, reconnect = True)

    def add_measurement(self, new_measurement: Measurement):
        """
        Adds a class which inherits `Measurement` to the program and adds it
        as a QCoDeS sub-module
        
        Args:
            new_measurement (Measurement): The instance which inherits
            Measurement to be added
        """
        self._measurements.append(new_measurement)

    def run(self, qua_program, **kwargs):
        """
        Sends the qua program for execution to the OPX and sets the programs 
        result handles 
        
        Args:
            qua_program (program): QUA program to be executed
        """
        self.qm_job = self.opx.execute(qua_program, **kwargs)
        self.result_handles = self.qm_job.result_handles

    def print_qua_program_to_file(
            self,
            path: str,
            qua_program,
            add_config: bool = False
            ) -> None:
        """
        Creates file with 'filename' and prints the QUA code to this file
        
        Args:
            file_name (str): File name of target file
            qua_program (program): QUA program to be printed
            add_config (bool): Whether config is added to output file
        """
        with open(path, 'w', encoding="utf-8") as file:
            if self.device is not None and add_config:
                file.write(generate_qua_script(
                    qua_program, self.device.config
                    ))
            else:
                file.write(generate_qua_script(qua_program))

    def run_local_simulation(
            self,
            qua_program,
            duration: int,
            **kwargs
            ):
        """
        Simulates the given program of the sequence for `duration` cycles
        TODO: Move to SequenceBase and add checks if OPX is connected
        Args:
            qua_program (program): QUA program to be simulated
            duration (int): Simulation duration in cycles
            **kwargs: Arbitrary keyword arguments for QMM simulation

        Returns:
            simulated_job (SimulatedJob): QM job with waveform simulation result
            nr_controllers (int): Nr of controllers to fetch simulation results
        """
        if not self.qmm:
            raise ConnectionError(
                "No QMM found! Connect an OPX via `connect_OPX`")
        simulated_job = self.qmm.simulate(
            self.opx.get_config(),
            qua_program,
            SimulationConfig(duration=duration),
            **kwargs
        )
        return simulated_job

    def get_idn(self):
        """
        Overload the get_idn method as we don't have one.
        """
        return None

    def ask_raw(self, cmd: str) -> str:
        """Abstract method from qcodes Instrument"""
        raise NotImplementedError

    def write_raw(self, cmd: str) -> None:
        """Abstract method from qcodes Instrument"""
        raise NotImplementedError

    def create_measurement_from_experiment(
            self,
            experiment: Experiment,
            qc_measurement_name: str | None = None,
            name: str = 'measurement',
            ) -> Measurement:
        """
        Creates an arbok and QCoDeS measurement from an arbok-experiment
        
        Args:
            name (str): Name of the measurement (py. variable name compliant)
            experiment (arbok_driver.Experiment): Experiment to be run
            qc_measurement_name (str): Name of the QCoDeS measurement
                (as it will be saved in the database)

        Returns:
            measurement (arbok_driver.Measurement): Measurement instance
        """
        measurement = Measurement(
            parent = self,
            name = name,
            sequence_config = self.device.param_config
            )
        if qc_measurement_name is None:
            qc_measurement_name = experiment.name
        measurement.qc_experiment = load_or_create_experiment(
            experiment.name, self.device.name)
        measurement.qc_measurement_name = qc_measurement_name
        sub_sequences = experiment.get_sequences_config(self.device)
        measurement.add_subsequences_from_dict(sub_sequences)
        return measurement

    def check_db_engine_and_bucket_connected(self):
        """
        Checks if database engine and s3 bucket are connected
        Raises error if not connected
        TODO: ping both TCP connections to check if still alive
        """
        if self.database_engine is None:
            raise ConnectionError(
                "No database engine connected! Please connect a database "
                "engine to the arbok_driver before running measurements.")
        if self.minio_filesystem is None:
            raise ConnectionError(
                "No MinIO filesystem connected! Please connect a MinIO filesystem to the "
                "arbok_driver before running measurements.")

    def get_run_from_id(self, run_id: int) -> SqlRun:
        """
        Fetches a run from the connected arbok database engine based on the
        given run ID
        
        Args:
            run_id (int): ID of the run to be fetched
        Returns:
            run (arbok_driver.sqlalchemy_classes.SqlRun): detached run instance
        """
        if self.database_engine is None:
            raise ConnectionError(
                "No database engine connected! Please connect a database "
                "engine to the arbok_driver before fetching runs.")
        with Session(self.database_engine) as session:
            sql_run = session.get(SqlRun, run_id)
        if sql_run is None:
            raise LookupError(
                f"No SqlRun found for run-ID: {run_id}"
            )
        return sql_run

    def get_data_from_id(self, run_id: int) -> xr.Dataset:
        """
        Fetches data from the connected database engine based on the
        given run ID
        
        Args:
            run_id (int): ID of the run to fetch data from
        Returns:
            xr_dataset (xarray.Dataset): Lazy loaded xarray dataset. Will only
                load data when using .load() or .compute() methods!
        """
        sql_run = self.get_run_from_id(run_id)
        minio_name = f"{sql_run.run_id}_{sql_run.uuid}"
        store = self.minio_filesystem.get_mapper(f'dev/{minio_name}/data.zarr')
        xr_dataset = xr.open_zarr(store, consolidated = True)
        return xr_dataset
