import json
import struct
import threading
from itertools import chain

from pymoku import (log, MokuException, DeployException)
from pymoku.network import NetworkError
from pymoku.sources import LoopMux, SlotOutput
from pymoku.plugins import MokuPlugins
from pymoku.registers import CachableAccessControl
from pymoku.progress import progress

from contextlib import contextmanager, ExitStack


class Moku(MokuPlugins, CachableAccessControl):
    """
    Core class representing a connection to a physical Moku:Lab unit.

    This must always be created first. Once a :any:`Moku` object exists, it can
    be queried for running Slot or new Slots deployed to the device.
    """

    def __init__(self, socket_factory, control, fileserv):
        self._ctx = ExitStack()
        self.socket_factory = socket_factory
        self._fileserv = self._ctx.enter_context(fileserv)
        self._control = self._ctx.enter_context(control)
        CachableAccessControl.__init__(self, 'platform', self._control)
        self._slots = []
        self.notif = None
        self._cached_state = None
        self._cached = False
        self._sources = []
        self._outputs = []
        device, system = self._control.get_properties(['device', 'system'])
        self.fw_ver = int(device['device']['fw_version'])
        self.hw_ver = device['device']['hw_version']
        self.serial = int(device['device']['serial'])
        self.bootmode = system['system'].get('bootmode', 'normal')
        self.platform_sha = None
        self._slots_lock = threading.Lock()

    def attach(self):
        with progress() as prog:
            prog.set_message(f"Loading {self.name}")
            # self.notif = self._ctx.enter_context(NotificationLogger(self))
            # self.notif.subscribe('hardware.modified', self.hardware_modified)
            self.load_plugins_for_moku()
            prog.set_progress(0.2)
            with self.cached_state() as state:
                prog.set_progress(0.4)
                self._update_params(state)
                prog.set_progress(0.9)
                self._update_slots(state)

    def _update_params(self, state):
        self._update_coeffs(state)
        self.platform_id = state['platform'].get('ID')
        self.platform_sha = state['platform'].get('sha')
        self.num_slots = state['platform'].get('slots', 0)
        self.num_loops = state['platform'].get('loops', 0)
        self.num_inputs = state['platform'].get('inputs', 0)
        self.num_slot_outputs = state['platform'].get('slot_outputs', 0)
        self._sys_ref = state['clk'].get('_sys_ref', 312.5e6)  # HACK
        self.setup_routing()
        self.load_functions_for_platform()
        self.install_empties_for_platform()

    def _update_slots(self, state):
        self._slots = [None] * self.num_slots
        for slot in range(self.num_slots):
            log.debug(f'load slot {slot:d}')

            if f'instr{slot:d}' not in state:
                continue

            if (slot_id := state[f'instr{slot:d}']['ID']) is None:
                self._slots[slot] = None
                continue

            if slot_id == 'pymoku:empty':
                continue

            try:
                manifest = self.function_manifest(slot_id, state[f'instr{slot:d}']['sha'])
                plugin = self.plugin_for_function(slot_id)
                slot_inst = plugin.attach(self, manifest)
                slot_inst.ID = slot_id
                slot_inst.PLUGIN = plugin
                self.attach_slot(slot_inst, slot)  # TODO maybe in the plugin?
                self._slots[slot] = slot_inst
            except Exception as e:
                log.exception(str(e))
                log.warning(f'Failed to load slot ID: {slot_id} in slot {slot:d}')

    @property
    def ip_addr(self):
        return self.socket_factory.ip_addr

    @property
    def name(self):
        return self._control.get_property('system.name')

    @name.setter
    def name(self, name):
        self._control.set_property('system.name', name)

    def deploy_platform(self, platform_id, plat_sha=None, force=False):
        with progress() as prog:
            prog.set_message(f"Deploying {platform_id}")
            self.install_barfile(self.platform_manifest(plat_sha))
            prog.set_progress(0.1)
            with self.cached_state(platform=dict(ID=platform_id,
                                   sha=plat_sha, force=force)) as state:
                prog.set_progress(0.7)
                self._update_params(state)
                prog.set_progress(0.9)
                self._update_slots(state)

    @contextmanager
    def _claim_slot(self, slot=None):
        with self._slots_lock:
            if slot is not None:
                yield slot
                return

            for slot_idx in range(len(self._slots)):
                slot_inst = self._slots[slot_idx]
                if slot_inst is None or slot_inst.ID == 'pymoku:empty':
                    yield slot_idx
                    return
            else:
                raise DeployException('No free slot')

    def deploy(self, _id, slot=None, force=0):
        with progress() as prog:
            prog.set_message(f'Deploy {_id}')
            prog.set_progress(0.1)

            with self._claim_slot(slot) as slot_idx:
                log.info(f"Deploy {_id} in slot {slot_idx:d}")
                manifest = self.function_for_slot(_id, self.platform_sha, slot_idx)
                assert manifest is not None, f"No barfile available for {_id}"
                plugin = self.plugin_for_function(_id)
                prog.set_progress(0.2)
                self.install_barfile(manifest)
                self.install_barfile(self.plugin_manifest(_id))
                prog.set_progress(0.5)

                plat_defn = {f'instr{slot_idx:d}': dict(ID=_id, force=bool(force))}
                with self.cached_state(**plat_defn):
                    prog.set_progress(0.6)
                    slot_inst = plugin.attach(self, manifest)
                    slot_inst.ID = _id
                    slot_inst.PLUGIN = plugin
                    prog.set_progress(0.8)
                    slot_inst.attach()
                    prog.set_progress(0.9)
                    slot_inst.reset()
                    self.attach_slot(slot_inst, slot_idx)

                self._slots[slot_idx] = slot_inst

        return slot_inst

    def _create_slot_instance(self, slot):
        with self.cached_state() as state:
            _id = state[f'instr{slot:d}']['ID']
            plugin = self.plugin_for_function(_id)
            manifest = self.function_manifest(_id, state[f'instr{slot:d}']['sha'])
            slot_inst = plugin.attach(self, manifest)
            slot_inst.ID = _id
            slot_inst.PLUGIN = plugin
            self._slots[slot] = slot_inst
            slot_inst.attach()
            self.attach_slot(slot_inst, slot)
            return slot_inst

    def clear_slot(self, slot):
        """
        Deploy the 'empty' bitstream to the slot'
        """
        self.deploy('pymoku:empty', slot, force=True)

    def sys_ref(self):
        return self._sys_ref

    def adc_sample_rate(self):
        return self.sys_ref() * 4.0

    def setup_routing(self):
        self._slot_outputs = [[SlotOutput(slot_idx, outp_idx)
                              for outp_idx in range(self.num_slot_outputs)]
                              for slot_idx in range(self.num_slots)]

        self._loop_muxes_1 = self._create_loop_muxes(width=1)
        self._loop_muxes_2 = self._create_loop_muxes(width=2)
        self._loop_muxes_4 = self._create_loop_muxes(width=4)

        for width in [1, 2, 4]:
            for plat_outp in self.platform_outputs():
                for idx, outp in enumerate(list(chain(*self._slot_outputs))[::width]):
                    edge_idx = tuple(range(idx * width, (idx + 1) * width, 1))
                    plat_outp.add(node=outp, idx=edge_idx, bus_width=width)

    def _create_loop_muxes(self, width):
        loop_muxes = [[LoopMux(slot_ctrl, loop_idx, width)
                      for loop_idx in range(0, self.num_loops, width)]
                      for slot_ctrl in self.routing.slots]

        for loop in chain(*loop_muxes):
            for idx, outp in enumerate(list(chain(*self._slot_outputs))[::width]):
                edge_idx = tuple(range(idx * width, (idx + 1) * width, 1))
                loop.add(node=outp, idx=edge_idx, bus_width=width)

        return loop_muxes

    def platform_sources(self):
        """
        Return the fixed platform sources (ADCs, DIO, etc).
        """
        return self._sources

    def slot_sources(self):
        """
        Return all currently available output sources from deployed slots.
        """
        return self._slot_outputs

    def sources(self):
        """
        Return all available sources in this platform, both fixed and slot.
        """
        return self.platform_sources() + self.slot_sources()

    def get_source(self, name):
        """
        Find and return the Source with the given name
        """
        for src in self.sources():
            if name == src.name:
                return src
        raise Exception(f'No such source: {name}')

    def platform_outputs(self):
        """
        Return the fixed outputs of this platform (DACs, DIO, etc).
        """
        return self._outputs

    def update_firmware(self, fw):
        """ Updates the firmware on the Moku.

        The Moku will automatically power off when the update is complete.
        """
        log.info(f"Uploading custom firmware: {fw}")
        with open(fw, 'rb') as f:
            data = f.read()

        try:
            self._fileserv.send('f', 'moku.fw', data)
            reply = self._control._transfer(bytearray([0x52, 0x01]))
            if reply[1]:
                raise MokuException(
                    "Firmware update failure {:d}".format(reply[1]))
        except NetworkError:
            # Sometimes the network connection goes down before the ack can be
            # received
            pass

    def _restart_board(self):
        self._control._transfer(bytearray([0x52, 0x02]))

    def modify_hardware(self, **kwargs):
        if self._cached and not kwargs:
            return self._cached_state

        data = b'\x54' + json.dumps(kwargs).encode('utf8')
        resp = self._control._transfer(data, timeout=10000)
        ptype, status = struct.unpack('<BB', resp[:2])
        if not status:
            state = json.loads(resp[2:].decode('utf8'))
            if self._cached:
                self._cached_state = state  # TODO should recurse update
            return state

        if status == 33:
            raise DeployException('Missing Bitstream')
        raise DeployException(f'Deploy error {status:d}')

    def hardware_modified(self, evt_idx, evt_id, data):
        self._update_coeffs(data['state'])
        for i in self._slots:
            if i is not None:
                i.hardware_modified(data['params'], data['state'])

    @contextmanager
    def cached_state(self, **kwargs):
        """ Cache hw state
            Only read cache, writes will still go to network,
            and will update local cache
        """
        c = self._cached  # local copy to maintain reentrant scope
        if not c:
            self._cached_state = self.modify_hardware(**kwargs)
            self._cached = True

        try:
            yield self._cached_state
        finally:
            if not c:
                self._cached = False

        if not c:
            self._cached_state = None

    def config(self):
        return self._control.get_property('')

    def bitstreams(self):
        return self._control.reload_bitstreams()

    def dac_source(self, dac_ch, source=None):
        dac = f'dac{dac_ch:d}'
        data = {dac: dict(source=source)} if source is not None else {}
        resp = self.modify_hardware(**data)
        return resp[dac]['source']

    def loop_source(self, slot, idx, source_idx=None):
        instr = f'instr{slot:d}'
        source = f'source{idx:d}'
        data = {instr: {source: source_idx}} if source_idx is not None else {}
        resp = self.modify_hardware(**data)
        return resp[instr][source]

    def use_ext_clock(self, ext):
        self.modify_hardware(clk=dict(ext_enable=ext))

    def write_calibration(self, cal_data=None, write_eeprom=False):
        return self._control.write_calibration(cal_data=cal_data, write_eeprom=write_eeprom)

    def get_calibration(self):
        return self._control.get_calibration()

    def get_statistics(self):
        try:
            return self._control.get_statistics()
        except NetworkError:
            return None

    def sync(self, mask):
        self.modify_hardware(instr_sync=mask)

    def cbufs(self):
        state = self.modify_hardware()
        slot_cbufs = state['platform']['slot_cbufs']
        slots = state['platform']['slots']
        return [[f'cbuf{cbuf + slot * slot_cbufs:d}'
                for cbuf in range(slot_cbufs)]
                for slot in range(slots)]

    def close(self):
        with self._ctx:
            for slot_inst in self._slots:
                if slot_inst and slot_inst.attached:
                    slot_inst.dettach()
            self._control.ownership(release=None)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    @staticmethod
    def connect(addr):
        """ Like pymoku.connect() but without all the version checking
        """
        from pymoku import ZMQCurveSocketFactory, Control, FileTransfer
        skt_factory = ZMQCurveSocketFactory(addr)
        control = Control(skt_factory.getControlSocket())
        fileserv = FileTransfer(control.skt)
        moku = MOKU_CLASS[control.get_property('device.hw_version')](
            skt_factory, control, fileserv)
        return moku


from pymoku.platforms.mokulab import MokuLab  # noqa
from pymoku.platforms.mokugo import MokuGo  # noqa
from pymoku.platforms.mokupro import MokuPro  # noqa
from pymoku.platforms.mokudelta import MokuDelta  # noqa


MOKU_CLASS = {  # Class for hwver
    '0.9': MokuLab,
    '1.0': MokuLab,
    '2.0': MokuLab,
    '3.0': MokuPro,
    '4.0': MokuGo,
    # '5.0': ZCU208,
    '6.0': MokuDelta,
}
