from collections import OrderedDict
from multiprocessing.connection import Connection

import numpy as np

from urh.dev.native.Device import Device
from urh.dev.native.lib import sdrplay
from urh.util.Logger import logger


class SDRPlay(Device):
    DEVICE_LIB = sdrplay
    ASYNCHRONOUS = True
    DEVICE_METHODS = Device.DEVICE_METHODS.copy()
    DEVICE_METHODS[Device.Command.SET_RF_GAIN.name] = "set_gain"
    DEVICE_METHODS[Device.Command.SET_IF_GAIN.name]["rx"] = "set_if_gain"
    DEVICE_METHODS[Device.Command.SET_ANTENNA_INDEX.name] = "set_antenna"

    DATA_TYPE = np.int16

    def __init__(
        self,
        center_freq,
        sample_rate,
        bandwidth,
        gain,
        if_gain=1,
        baseband_gain=1,
        resume_on_full_receive_buffer=False,
    ):
        super().__init__(
            center_freq=center_freq,
            sample_rate=sample_rate,
            bandwidth=bandwidth,
            gain=gain,
            if_gain=if_gain,
            baseband_gain=baseband_gain,
            resume_on_full_receive_buffer=resume_on_full_receive_buffer,
        )
        self.success = 0
        self.error_codes = {
            0: "SUCCESS",
            1: "FAIL",
            2: "INVALID PARAMETER",
            3: "OUT OF RANGE",
            4: "GAIN UPDATE ERROR",
            5: "RF UPDATE ERROR",
            6: "FS UPDATE ERROR",
            7: "HARDWARE ERROR",
            8: "ALIASING ERROR",
            9: "ALREADY INITIALIZED",
            10: "NOT INITIALIZED",
            11: "NOT ENABLED",
            12: "HARDWARE VERSION ERROR",
            13: "OUT OF MEMORY ERROR",
        }

    @staticmethod
    def device_dict_to_string(d):
        hw_ver = d["hw_version"]
        serial = d["serial"]
        return "RSP {} ({})".format(hw_ver, serial)

    @property
    def device_parameters(self):
        return OrderedDict(
            [
                (self.Command.SET_ANTENNA_INDEX.name, self.antenna_index),
                (self.Command.SET_FREQUENCY.name, self.frequency),
                (self.Command.SET_SAMPLE_RATE.name, self.sample_rate),
                (self.Command.SET_BANDWIDTH.name, self.bandwidth),
                (self.Command.SET_RF_GAIN.name, self.gain),
                (self.Command.SET_IF_GAIN.name, self.if_gain),
                ("identifier", self.device_number),
            ]
        )

    @property
    def has_multi_device_support(self):
        return True

    @classmethod
    def get_device_list(cls):
        return [cls.device_dict_to_string(d) for d in sdrplay.get_devices()]

    @classmethod
    def enter_async_receive_mode(
        cls, data_connection: Connection, ctrl_connection: Connection
    ):
        ret = sdrplay.init_stream(
            cls.sdrplay_initial_gain,
            cls.sdrplay_initial_sample_rate,
            cls.sdrplay_initial_freq,
            cls.sdrplay_initial_bandwidth,
            cls.sdrplay_initial_if_gain,
            data_connection,
        )

        ctrl_connection.send(
            "Start RX MODE with \n  FREQUENCY={}\n  SAMPLE_RATE={}\n  BANDWIDTH={}\n  GAIN={}\n  IF_GAIN={}:{}".format(
                cls.sdrplay_initial_freq,
                cls.sdrplay_initial_sample_rate,
                cls.sdrplay_initial_bandwidth,
                cls.sdrplay_initial_gain,
                cls.sdrplay_initial_if_gain,
                ret,
            )
        )

        return ret

    @classmethod
    def init_device(
        cls, ctrl_connection: Connection, is_tx: bool, parameters: OrderedDict
    ) -> bool:
        identifier = parameters["identifier"]

        try:
            device_list = sdrplay.get_devices()
            device_number = int(identifier)
            ctrl_connection.send(
                "CONNECTED DEVICES: {}".format(
                    ", ".join(map(cls.device_dict_to_string, device_list))
                )
            )
            ret = sdrplay.set_device_index(device_number)
            ctrl_connection.send(
                "SET DEVICE NUMBER to {}:{}".format(device_number, ret)
            )
        except (TypeError, ValueError) as e:
            logger.exception(e)
            return False

        device_model = device_list[device_number]["hw_version"]
        sdrplay.set_gr_mode_for_dev_model(device_model)
        if device_model == 2:
            antenna = parameters[cls.Command.SET_ANTENNA_INDEX.name]
            cls.process_command(
                (cls.Command.SET_ANTENNA_INDEX.name, antenna),
                ctrl_connection,
                is_tx=False,
            )
        else:
            ctrl_connection.send("Skipping antenna selection for RSP1 device")

        cls.sdrplay_initial_freq = parameters[cls.Command.SET_FREQUENCY.name]
        cls.sdrplay_initial_sample_rate = parameters[cls.Command.SET_SAMPLE_RATE.name]
        cls.sdrplay_initial_bandwidth = parameters[cls.Command.SET_BANDWIDTH.name]
        cls.sdrplay_initial_gain = parameters[cls.Command.SET_RF_GAIN.name]
        cls.sdrplay_initial_if_gain = parameters[cls.Command.SET_IF_GAIN.name]
        cls.sdrplay_device_index = identifier
        return True

    @classmethod
    def shutdown_device(cls, ctrl_connection, is_tx: bool):
        logger.debug("SDRPLAY: closing device")
        ret = sdrplay.close_stream()
        ctrl_connection.send("CLOSE STREAM:" + str(ret))

        if cls.sdrplay_device_index is not None:
            ret = sdrplay.release_device_index()
            ctrl_connection.send("RELEASE DEVICE:" + str(ret))

    @staticmethod
    def bytes_to_iq(buffer):
        return np.frombuffer(buffer, dtype=np.int16).reshape((-1, 2), order="C")
