import copy
import pickle
import threading
import time
import numpy as np
import zmq
from typing import Iterable
import blosc2
from typing import List
from npy_append_array import NpyAppendArray
from logging import getLogger
from datetime import datetime
import os
from hashlib import md5
from mkidgen3.server.misc import zpipe
from mkidgen3.util import format_bytes
from mkidgen3.funcs import convert_adc_raw_to_mv, raw_iq_to_unit, raw_phase_to_radian
from ..system_parameters import SYSTEM_BANDWIDTH
from mkidgen3.rfsocmemory import memfree_mib
from .feedline_config import FeedlineConfig, WaveformConfig, IFConfig
from ..mkidpynq import PHOTON_DTYPE
from ..system_parameters import N_CHANNELS, N_IQ_GROUPS, N_PHASE_GROUPS, N_POSTAGE_CHANNELS, \
    PHOTON_POSTAGE_WINDOW_LENGTH, MAXIMUM_DESIGN_COUNTRATE_PER_S
from functools import cached_property
from mkidgen3.mkidpynq import unpack_photons

from typing import Tuple
PHOTON_DTYPE_PACKED = np.uint64


class FRSClient:
    def __init__(self, url: str, command_port: int = 10000, data_port: int = 10001, status_port: int = 10002):
        """
        Args:
            url: ip address of rfsoc for ex: rfsoc4x2b.physics.ucsb.edu
            command_port: port for feedline server commands
            data_port: port for capture data server
            status_port: port for status server
        """
        self.url = url
        self.command_port = command_port
        self.data_port = data_port
        self.status_port = status_port

    @property
    def command_url(self):
        return f'tcp://{self.url}:{self.command_port}'

    @property
    def data_url(self):
        return f'tcp://{self.url}:{self.data_port}'

    @property
    def status_url(self):
        return f'tcp://{self.url}:{self.status_port}'

    def __str__(self):
        return f'FRS@{self.url}:{self.command_port}'

    def bequiet(self):
        ctx = zmq.Context().instance()
        with ctx.socket(zmq.REQ) as s:
            s.connect(self.command_url)
            s.send_pyobj(('bequiet', {}))
            return s.recv_pyobj()

    def order(self, *args):
        ctx = zmq.Context().instance()
        with ctx.socket(zmq.REQ) as s:
            s.connect(self.command_url)
            s.send_pyobj(args)
            return s.recv_pyobj()


class CaptureRequest:
    """
    A Capture Request object is generated by the client and submitted to the feedline readout server
    """
    STATUS_ENDPOINT = 'inproc://cap_stat.xsub'
    DATA_ENDPOINT = 'inproc://cap_data.xsub'
    FINAL_STATUSES = ('finished', 'aborted', 'failed')
    SUPPORTED_TAPS = ('postage', 'photon', 'adc', 'iq', 'filtphase', 'ddciq')

    @staticmethod
    def validate_channels(tap: str, chan: list):
        """
        Verify that the channel selection is compatible with the driver.

        Note that this is not done dynamically, no driver discovery is performed

        See mkidgen3.drivers.FilterPhase, .drivers.FilterIQ, and .trigger.PhotonPostageFilter
        for details about groups and monitor channels.

        Args:
            chan: A list of channels
            tap: A capture location

        Returns: True|False
        """
        tap = tap.lower()
        if 'iq' in tap:
            l = N_CHANNELS
            m = N_CHANNELS
        elif 'phase' in tap:
            l = N_CHANNELS
            m = N_CHANNELS
        elif tap == 'postage':
            l = N_POSTAGE_CHANNELS
            m = N_CHANNELS
        elif chan:
            l = 0  # no channels for photon or adc capture
        else:
            return False
        try:
            assert len(chan) <= l
            for c in chan:
                assert 0 <= c < m
        except AssertionError:
            return False
        return True

    def __init__(self, n: int, tap: str, feedline_config: FeedlineConfig,
                 feedline_server: FRSClient, channels: Iterable | int | None = None, file: str = None,
                 fail_saturation_frac: None | float | int = None, numpy_metric: None | str = None,
                 _compression_override: bool = None):
        """

        Args:
            n: the number of samples to capture, for photons n is the buffer rotation interval request in ms
            tap: the location to capture: 'adc'|'iq'|'phase'|'postage'|'photon'
            feedline_config: the FL configuration required by the capture
            feedline_server: the FRS to capture from
            channels: the phase, IQ, or postage resonator channels to capture
            file: an optional file to save the result to
            fail_saturation_frac: Take a pre-capture of ADC data and abort if the maximum I or Q magnitude is greater
                                  than the fraction*ADC_MAX_INT. Set to None, 0, or > 1 to disable.
            numpy_metric: a numpy function to apply across samples on a channel by channel basis. Must have `out` and
                          `axis` as input parameters. Math is done to I and Q respectively (complex math not supported).
                          For example: 'mean' will return a mean of the values for each channel for IQ and phase data.
            _compression_override: skips blosc compression on the server if True. False or None (default) results in
                                   data being compressed before sending.
        """

        tap = tap.lower()
        assert tap in CaptureRequest.SUPPORTED_TAPS
        self.fail_saturation_frac = float(fail_saturation_frac) if fail_saturation_frac else None
        self.nsamp = n  # n is treated as the buffer time in ms for photons, and has limits enforced by the driver
        self._last_status = None
        if channels is not None:
            try:
                channels = tuple(sorted(set(channels)))
            except TypeError:
                channels = (channels,)
            if not CaptureRequest.validate_channels(tap, channels):
                raise ValueError('Invalid channels specification')
        self.channels = channels
        self.tap = tap  # maybe add some error handling here
        self.feedline_config = feedline_config
        self.server = feedline_server
        self._status_socket = None
        self._data_socket = None
        self._established = False
        self._compression_override = _compression_override
        if numpy_metric:
            try:
                self.numpy_metric = str(numpy_metric)
                getattr(np, self.numpy_metric)
            except AttributeError:
                raise ValueError('Not a good metric')
        else:
            self.numpy_metric = None
        self.data_endpoint = file or type(self).DATA_ENDPOINT

    def __hash__(self):
        return int(md5(str((hash(self.feedline_config), self.tap, self.data_endpoint,
                            self.nsamp, self.server.command_url)).encode()).hexdigest(), 16)

    # def __enter__(self):
    #     self.establish()
    #
    # def __exit__(self, exc_type, exc_val, exc_tb):
    #     if isinstance(exc_type, zmq.error.ContextTerminated):
    #         self._data_socket=None
    #         self._status_socket=None
    #     else:
    #         self.destablish()

    def __del__(self):
        try:
            self.destablish()
        except zmq.error.ContextTerminated:
            pass

    def __str__(self):
        return f'CapReq {self.server}:{self.tap} {str(hash(self))}'

    def submit(self, ctx: zmq.Context = None):
        getLogger(__name__).info(f'Submitting {self}..')
        ctx = ctx or zmq.Context().instance()
        with ctx.socket(zmq.REQ) as s:
            s.connect(self.server.command_url)
            s.send_pyobj(('capture', self))
            getLogger(__name__).info(f'...submitted...')
            resp = s.recv_pyobj()
            getLogger(__name__).info(f'response: {resp}')
            return resp

    @property
    def type(self):
        return 'engineering' if self.tap not in ('photon', 'postage') else self.tap

    @property
    def id(self):
        return str(hash(self)).encode()

    def establish(self, context: zmq.Context = None):
        """
        Open up the outbound sockets for status and data. If data is going to a file, verify it ca be opened for writing and close it
        """
        if self._status_socket is not None or self._data_socket is not None:
            getLogger(__name__).warning(f'{self} already established')
            return
        context = context or zmq.Context.instance()
        self._status_socket = context.socket(zmq.PUB)
        self._status_socket.connect(self.STATUS_ENDPOINT)
        if self.data_endpoint.startswith('file://'):
            f = self.data_endpoint.lstrip('file://')
            if not os.access('/' + os.path.dirname(f), os.W_OK):
                raise IOError(f'Unable to write to {f}')
            return
        self._data_socket = context.socket(zmq.PUB)
        self._data_socket.connect(self.data_endpoint)
        self._send_status('established')

    def destablish(self):
        if self._status_socket is not None:
            log = True
            getLogger(__name__).debug(f'de-establishing status...')
        try:
            # NB the context or socket's linger setting determines how long any pending messages have to get sent off
            self._status_socket.close()
            self._status_socket = None
        except AttributeError:
            pass
        if self._data_socket is not None:
            log = True
            getLogger(__name__).debug(f'de-establishing data...')
        try:
            self._data_socket.close()
            self._data_socket = None
        except AttributeError:
            pass
        if log:
            getLogger(__name__).debug(f'...de-established.')

    def fail(self, message, raise_exception=False):
        if self.completed:
            return
        try:
            if self.data_endpoint.startswith('file://'):
                pass
            else:
                self._data_socket.send_multipart([self.id, b''])
            self._send_status('failed', message)
            self.destablish()
        except zmq.ZMQError as ez:
            getLogger(__name__).warning(f'Failed to send fail message {self} due to {ez}')
            if raise_exception:
                raise ez

    def finish(self, raise_exception=True):
        if self.completed:
            return
        try:
            if self.data_endpoint.startswith('file://'):
                pass
            else:
                self._data_socket.send_multipart([self.id, b''])
            self._send_status('finished', f'Finished at UTC {datetime.utcnow()}')
            self.destablish()
        except zmq.ZMQError as ez:
            getLogger(__name__).warning(f'Failed to send finished message {self} due to {ez}')
            if raise_exception:
                raise ez

    def abort(self, message, raise_exception=False):
        if self.completed:
            return
        try:
            self._data_socket.send_multipart([self.id, b''])
            self._send_status('aborted', message)
            self.destablish()
        except zmq.ZMQError as ez:
            getLogger(__name__).warning(f'Failed to send abort message {self} due to {ez}')
            if raise_exception:
                raise ez

    def send_data(self, data, status='', copy=False, compress=True):
        """
        Send a (chunk) of data out. By default, data is published as a blosc2 compressed array with the
        capture id as the subscription key. CaptureRequests created with a file destination  will be saved as
        uncompressed npy data, appended on successive calls, no data will be published.

        Data is shipped in the (compressed) format it is given.

        Do not call this function from multiple threads.
        Call .establish() in the calling thread before calling this function or after calling destablish()

        This function is intended to be zero-copy, is as of 10/27/23 and should be maintained as such.

        Args:
            data: a array of data to send out
            status: (optional) a status update message to send out
            copy: copy the array first, you probably don't want this. copying the array then freeing the
            underlying buffer causes segfaults with the pycharm debugger. Be aware that the property
            zmq.COPY_THRESHOLD will case messages to be copied even if copy is False if the message size
            is less than the threshold.
            compress: compress the data with blosc2, ignored for file destinations.

        Returns: None | zmq.MessageTracker None is returned if copying or the destination is a file. A
        segfault may occur if the data is deallocated prior to the messagetracker being done.
        """
        if not self._status_socket:
            raise RuntimeError('Status socket is not established.')
        if self.data_endpoint.startswith('file://'):
            with NpyAppendArray('/' + self.data_endpoint.lstrip('file://'), delete_if_exists=False) as x:
                x.append(data)
            self._send_status('capturing', status)
            return
        elif not self._data_socket:
            raise RuntimeError('Data socket is not established.')

        #        getLogger(__name__).debug(f'MiB Free: {memfree_mib()}')
        if self.numpy_metric is not None:
            out = np.empty(data.shape[1:] if data.ndim >1 else (1,))
            try:
                getattr(np, self.numpy_metric)(data, axis=0, out=out)
            except TypeError as e:
                raise RuntimeError(f'numpy metric: {self.numpy_metric} failed with {e}')
            except Exception as e:
                raise RuntimeError(f'Encountered unexpected error related to numpy metric {e}')
            data = out
        else:
            data = np.array(data) if copy else data

        do_compression = compress if self._compression_override is None else not self._compression_override

        # from line_profiler import LineProfiler
        # profile = LineProfiler()
        # profile.enable_by_count()
        times = []
        times.append(time.perf_counter())
        #        getLogger(__name__).debug(f'MiB Free: {memfree_mib()}')
        compressed = blosc2.compress(data) if do_compression else data
        times.append(time.perf_counter())
        #        getLogger(__name__).debug(f'MiB Free: {memfree_mib()}')
        self._send_status('capturing', status)
        times.append(time.perf_counter())
        tracker = self._data_socket.send_multipart([self.id, compressed], copy=False, track=not copy)
        times.append(time.perf_counter())
        # profile.disable_by_count()
        #getLogger(__name__).debug(profile.get_stats())
        times = np.diff(times) * 1000
        getLogger(__name__).debug(f'Compress {times[0]:.2f} ms, Status Up {times[1]:.2f}, Send {times[2]:.2f}')
        cval = 100 * len(compressed) / data.nbytes if data.nbytes else 100
        getLogger(__name__).debug(f'Sending {format_bytes(len(compressed))}, '
                                  f'compressed to {cval:.1f} % from {format_bytes(data.nbytes)}.')
        return tracker

    def _send_status(self, status, message=''):
        if not self._status_socket:
            raise RuntimeError('No status_socket connection available')
        update = f'{status}:{message}'
        getLogger(__name__).getChild('statusio').debug(f'Published status update for {self}: "{update}"')
        self._last_status = (status, message)
        self._status_socket.send_multipart([self.id, update.encode()])

    @property
    def completed(self):
        return self._last_status is not None and self._last_status[0] in ('finished', 'aborted', 'failed')

    def set_status(self, status, message='', context: zmq.Context = None, destablish=False):
        """
        get appropriate context and send current status message after connecting the status socket

        context is ignored if a status destination connection is extant
        """
        if not self._status_socket:
            context = context or zmq.Context().instance()
            self._status_socket = context.socket(zmq.PUB)
            self._status_socket.connect(self.STATUS_ENDPOINT)
        self._send_status(status, message)
        if destablish:
            self.destablish()

    @property
    def size_bytes(self):
        if self.numpy_metric:
            return self.capture_atom_bytes
        else:
            return self.nsamp*self.capture_atom_bytes

    @property
    def nchan(self):
        if self.channels:
            return len(self.channels)
        elif self.tap == 'postage':
            return N_POSTAGE_CHANNELS
        elif 'iq' in self.tap or 'phase' in self.tap:
            return N_CHANNELS
        else:
            return 1

    @property
    def dwid(self) -> int:
        """Data size of a capture sample in bytes"""
        if self.numpy_metric:
            return 2 * 8  ## numpy math returns float64 (8 bytes) * 2 for IQ
        if 'iq' in self.tap or self.tap in ('adc', 'postage'):
            return 4
        elif 'phase' in self.tap:
            return 2
        else:
            return PHOTON_DTYPE.itemsize

    @property
    def capture_atom_bytes(self) -> int:
        return self.dwid*self.nchan

    @property
    def buffer_shape(self):
        """ (nsamples, nchannels, 1|2)  """
        n = int(np.ceil(self.nsamp * MAXIMUM_DESIGN_COUNTRATE_PER_S / 1000)) if self.tap == 'photon' else self.nsamp

        if self.tap == 'postage':
            return n, PHOTON_POSTAGE_WINDOW_LENGTH + 1, 2
        elif 'iq' in self.tap or self.tap == 'adc':
            if self.numpy_metric:
                return self.nchan, 2
            else:
                return n, self.nchan, 2
        elif self.tap == 'photon':
            return (n,)
        elif 'phase' in self.tap:
            if self.numpy_metric:
                return self.nchan
            else:
                return n, self.nchan
        else:
            raise RuntimeError('Unknown tap: {self.tap}')


class CaptureSink(threading.Thread):
    def __init__(self, req_nfo: [CaptureRequest, tuple], start=True):
        if isinstance(req_nfo, CaptureRequest):
            id, buffer_shape, size_bytes, source_url, _compression_override = req_nfo.id, req_nfo.buffer_shape, req_nfo.size_bytes, req_nfo.server.data_url, req_nfo._compression_override
        else:
            id, buffer_shape, size_bytes, source_url, _compression_override = req_nfo
        super(CaptureSink, self).__init__(name=f'cap_id={id.decode()}')
        self._expected_buffer_shape = buffer_shape
        self._expected_bytes = size_bytes
        self.daemon = True
        self.cap_id = id
        self.data_source = source_url
        self.result = None
        self._pipe = None
        self._buf = []
        self._compression_override = _compression_override
        self.socket = None
        if start:
            self.start()

    def start(self):
        assert not self._started.is_set()
        self._pipe = zpipe(zmq.Context.instance())
        super(CaptureSink, self).start()

    def kill(self):
        self._pipe[0].send(b'')
        self._pipe[0].recv()
        self._pipe[1].close()
        self._pipe[0].close()

    def _accumulate_data(self, d):
        self._buf.append(d)

    def _finish_accumulation(self):
        self._buf = b''.join(self._buf)

    def _finalize_data(self):
        raise NotImplementedError

    def establish(self, ctx: zmq.Context = None):
        ctx = ctx or zmq.Context.instance()
        self.socket = ctx.socket(zmq.SUB)
        self.socket.setsockopt(zmq.SUBSCRIBE, self.cap_id)
        self.socket.connect(self.data_source)
        return self.socket

    def receive(self):
        id, data = self.socket.recv_multipart(copy=False)
        if not data:
            self.socket.close()
            return None
        return blosc2.decompress(data) if not self._compression_override else data

    def destablish(self):
        self.socket.close()
        self.socket = None

    def run(self):
        assert self.socket is None
        try:
            self.establish()
            poller = zmq.Poller()
            poller.register(self._pipe[1], flags=zmq.POLLIN)
            poller.register(self.socket, flags=zmq.POLLIN)

            getLogger(__name__).debug(f'Listening for data for {self.cap_id}')
            self._pipe[1].send(b'')
            while True:
                avail = dict(poller.poll())
                if self._pipe[1] in avail:
                    getLogger(__name__).debug(f'Received shutdown order, terminating data acq. of {self}')
                    break
                data = self.receive()  # does decompression
                if not data:
                    getLogger(__name__).debug(f'Received null, capture data stream over')
                    break
                getLogger(__name__).debug(f'Received {format_bytes(len(data))} of {format_bytes(self._expected_bytes)} for {self}')
                self._accumulate_data(data)
            self._finish_accumulation()
            self._finalize_data()
            if self.result:
                getLogger(__name__).info(f'Capture data for {self.cap_id} processed into {self.result.data.shape} '
                                         f'{self.result.data.dtype}: {self.result}')
            else:
                getLogger(__name__).info(f"Capture data for {self.cap_id} processed 'None'")

        except zmq.ZMQError as e:
            getLogger(__name__).warning(f'Shutting down {self} due to {e}')
        except AttributeError:
            raise
        finally:
            self.destablish()
            self._pipe[1].send(b'')

    def data(self, timeout=None):
        self.join(timeout=timeout)
        return self.result

    def ready(self):
        self._pipe[0].recv()  # block until a byte arrives


class StreamCaptureSink(CaptureSink):
    def _finalize_data(self):
        # raw adc data is i0q0 i1q1 int16
        bytes_recv = len(self._buf)
        if not bytes_recv:
            getLogger(__name__).warning(f'No data received for {self.cap_id}. '
                                        f'Expected {format_bytes(self._expected_bytes)}.')
            return
        elif bytes_recv > self._expected_bytes:
            getLogger(__name__).warning(f'Received more data than expected for {self.cap_id}. '
                                        f'Expected {format_bytes(self._expected_bytes)} got '
                                        f'{format_bytes(bytes_recv)}.')
        elif bytes_recv < self._expected_bytes:
            getLogger(__name__).warning(f'Finalizing incomplete capture data for {self.cap_id}. '
                                        f'Expected {format_bytes(self._expected_bytes)} got '
                                        f'{format_bytes(bytes_recv)}.')


        samples_recv = int(len(self._buf) / 2 // np.prod(self._expected_buffer_shape[1:]))  # / 2 for 2 bytes per I or Q
        shape = (min(samples_recv, self._expected_buffer_shape[0]),) + self._expected_buffer_shape[1:]

        dtype_bytes = self._expected_bytes // np.prod(shape, dtype=int)
        if dtype_bytes == 8:
            dtype = np.float64
        else:
            dtype = np.int16
            if dtype_bytes != 2:
                getLogger(__name__).warning(f'Received data has {dtype_bytes} bytes per sample. '
                                            f'Data may be lost in cast to int16 (2 bytes).')

        # TODO this might technically be a memory leak if we captured more data
        self.result = np.frombuffer(self._buf, count=np.prod(shape, dtype=int), dtype=dtype).reshape(shape).squeeze()


class ADCCaptureSink(StreamCaptureSink):
    def _finalize_data(self):
        super()._finalize_data()
        if self.result is not None:
            self.result = ADCCaptureData(self.result)


class IQCaptureSink(StreamCaptureSink):
    def _finalize_data(self):
        super()._finalize_data()
        if self.result is not None:
            self.result = IQCaptureData(self.result)


class PhaseCaptureSink(StreamCaptureSink):
    def _finalize_data(self):
        super()._finalize_data()
        if self.result is not None:
            self.result = PhaseCaptureData(self.result)


class SimplePhotonSink(CaptureSink):
    def _accumulate_data(self, d):
        self._buf.append(d)
        if sum(map(len, self._buf)) / 1024 ** 2 > 100:  # limit to 100MiB
            self._buf.pop(0)

    def _finish_accumulation(self):
        self._buf = b''.join(self._buf)

    def _finalize_data(self):
        self.result = PhotonCaptureData(self._buf)


class PhotonCaptureSink(CaptureSink):

    def capture(self, hdf, xymap, feedline_source, fl_ids):
        t = threading.Thread(target=self._main, args=(hdf, xymap, feedline_source, fl_ids))
        t.start()

    @staticmethod
    def _main(hdf, xymap, feedline_source, fl_ids, term_source='inproc://PhotonCaptureSink.terminator.inproc'):
        """

        Args:
            xymap: [nfeedline, npixel, 2] array
            feedline_source: zmq.PUB socket with photonbufers published by feedline
            term_source: a zmq socket of undecided type for detecting shutdown requests

        Returns: None

        """

        fl_npix = 2048
        n_fl = 5
        MAX_NEW_PHOTONS = 5000
        DETECTOR_SHAPE = (128, 80)
        fl_id_to_index = np.arange(n_fl, dtype=int)

        context = zmq.Context.instance()
        term = context.socket(zmq.SUB)
        term.setsockopt(zmq.SUBSCRIBE, id)
        term.connect(term_source)

        data = context.socket(zmq.SUB)
        data.setsockopt(zmq.SUBSCRIBE, fl_ids)
        data.connect(feedline_source)

        poller = zmq.Poller()
        poller.register(term, flags=zmq.POLLIN)
        poller.register(data, flags=zmq.POLLIN)

        live_image = np.zeros(DETECTOR_SHAPE)
        live_image_socket = None
        live_image_by_fl = live_image.reshape(n_fl, fl_npix)
        photons_rabuf = np.recarray(MAX_NEW_PHOTONS,
                                    dtype=(('time', 'u32'), ('x', 'u32'), ('y', 'u32'),
                                           ('phase', 'u16')))

        while True:
            avail = poller.poll()
            if term in avail:
                break

            frame = data.recv_multipart(copy=False)
            fl_id = frame[0]
            time_offset = frame[1]
            d = blosc2.decompress(frame[2])
            frame_duration = None  # todo time coverage of data
            # buffer is nchan*nmax+1 32bit: 16bit time(base2) 16bit phase
            # make array of to [nnmax+1, nchan, 2] uint16
            # nmax will always be <<2^12 number valid will be at [0,:,0]
            # times need oring w offset
            # photon data is d[1:d[0,i,0], i, :]

            nnew = d[0, :, 0].sum()
            # if we wanted to save binary data then we could save this, the x,y list, and the time offset
            # mean pixel count rate in this packet is simply [0,:,0]/dt
            fl_ndx = fl_id_to_index[fl_id]
            live_image_by_fl[fl_ndx, :] += d[0, :, 0] / frame_duration

            # if live_image_ready
            live_image_socket.send_multipart([f'liveim', blosc2.compress(live_image)])

            cphot = np.cumsum(d[0, :, 0], dtype=int)
            for i in range(fl_npix):
                sl_out = slice(cphot[i], cphot[i] + d[0, i, 0])
                sl_in = slice(1, d[0, i, 0])
                photons_rabuf['time'][sl_out] = d[sl_in, :, 0]
                photons_rabuf['time'][sl_out] |= time_offset
                photons_rabuf['phase'][sl_out] = d[sl_in, :, 1]
                photons_rabuf['x'][sl_out] = xymap[fl_ndx, i, 0]
                photons_rabuf['y'][sl_out] = xymap[fl_ndx, i, 1]
            hdf.grow_by(photons_rabuf[:nnew])

        term.close()
        data.close()
        hdf.close()


class PostageCaptureSink(CaptureSink):
    def _finalize_data(self):
        # get_postage returns:
        #  ids (nevents uint16 array or res cahnnels),
        #  events (nevents x 127 x2 of iq da

        if len(self._buf) != np.prod(self._expected_buffer_shape) * 4:
            getLogger(__name__).warning(f'Finalizing incomplete capture data for {self.cap_id}')
        self.result = PostageCaptureData(self._buf)


class StatusListener(threading.Thread):
    def __init__(self, id, source, initial_state='Created', start=True):
        super().__init__(name=f'StautsListner_{id.decode()}')
        self.daemon = True
        self.source = source
        self._pipe = None
        self.id = id
        self._status_messages = [initial_state]
        if start:
            self.start()

    def start(self):
        assert not self._started.is_set()
        self._pipe = zpipe(zmq.Context.instance())
        super(StatusListener, self).start()

    def kill(self):
        self._pipe[0].send(b'')
        self._pipe[0].recv()
        self._pipe[0].close()
        self._pipe[1].close()

    @staticmethod
    def is_final_status_message(update):
        """return True iff message is a final status update"""
        for r in CaptureRequest.FINAL_STATUSES:
            if update.startswith(r):
                return True
        return False

    def run(self):
        try:
            ctx = zmq.Context().instance()
            with ctx.socket(zmq.SUB) as sock:
                sock.linger = 0
                sock.setsockopt(zmq.SUBSCRIBE, self.id)
                sock.connect(self.source)

                poller = zmq.Poller()
                poller.register(self._pipe[1], flags=zmq.POLLIN)
                poller.register(sock, flags=zmq.POLLIN)
                getLogger(__name__).debug(f'Listening for status updates to {self.id}')
                self._pipe[1].send(b'')
                while True:
                    avail = dict(poller.poll())
                    if self._pipe[1] in avail:
                        getLogger(__name__).debug(f'Shutdown requested, terminating {self}')
                        break
                    elif sock not in avail:
                        time.sleep(.1)  # play nice
                        continue
                    id, update = sock.recv_multipart()
                    assert id == self.id or not self.id
                    update = update.decode()
                    self._status_messages.append(update)
                    getLogger(__name__).debug(f'Status update for {id}: {update}')
                    if self.is_final_status_message(update):
                        break
        except zmq.ZMQError as e:
            getLogger(__name__).critical(f"{self} died due to {e}")
        finally:
            self._pipe[1].send(b'')

    def latest(self):
        return self._status_messages[-1]

    def history(self):
        return tuple(self._status_messages)

    def ready(self):
        self._pipe[0].recv()  # TODO make this line block


class ADCCaptureData:
    """
    Formats raw ADC capture data to complex floats representing mV at the ADC input SMA.
    """

    def __init__(self, raw_data):
        """
        ADC data captured by an ADC capture request
        Args:
            raw_data: raw data captured by adc capture
        """
        self.raw = raw_data

    @property
    def dtype(self):
        return self.raw.dtype

    @property
    def shape(self):
        return self.raw.shape

    @cached_property
    def data(self):
        return convert_adc_raw_to_mv(self.raw[:, 0] + 1j * self.raw[:, 1])


class IQCaptureData:
    """
    Formats raw IQ capture data to complex floats between +/- 1.
    """

    def __init__(self, raw_data):
        self.raw = raw_data

    @property
    def dtype(self):
        return self.raw.dtype

    @property
    def shape(self):
        return self.raw.shape

    @cached_property
    def data(self):
        return raw_iq_to_unit(self.raw[..., 0] + self.raw[..., 1] * 1j)


class PhaseCaptureData:
    def __init__(self, data):
        self.raw = data

    @property
    def dtype(self):
        return self.raw.dtype

    @property
    def shape(self):
        return self.raw.shape

    @cached_property
    def data(self):
        return raw_phase_to_radian(self.raw, scaled=False)


class PhotonCaptureData:
    def __init__(self, buf):
        self.raw = np.frombuffer(buf, dtype=PHOTON_DTYPE_PACKED)
        self.photons = unpack_photons(self.raw)

    @property
    def data(self):
        return self.photons

    @property
    def dtype(self):
        return PHOTON_DTYPE

    @property
    def shape(self):
        return self.photons.size


class PostageCaptureData:
    def __init__(self, buf) :
        result = np.frombuffer(buf, count=len(buf) // 2, dtype=np.int16)
        n = result.size//2//128
        shape = (n, 128, 2)
        result = result.reshape(shape)
        from ..funcs import postage_buffer_to_data
        ids, events = postage_buffer_to_data(result, complex=True, scaled=True)
        self.ids = ids  # result[:, 0, 0].astype(np.uint16)
        self.iqs = events  # result[:, 1:, :]
        self.raw = result

    @property
    def dtype(self):
        return np.int16

    @property
    def shape(self):
        return self.iqs.shape

    @property
    def data(self):
        return self.iqs


class CaptureJob:
    def __init__(self, request: CaptureRequest, submit=False):
        """
        Args:
            request:
            submit: tuple of bools (with sink?, with status?) or True or False
        """
        self.request = request
        self.submitted = False

        self._status_listener = StatusListener(request.id, request.server.status_url,
                                               initial_state='created', start=False)

        if request.tap == 'adc':
            datasink = ADCCaptureSink
        elif 'iq' in request.tap:
            datasink = IQCaptureSink
        elif 'phase' in request.tap:
            datasink = PhaseCaptureSink
        elif request.tap == 'photon':
            datasink = SimplePhotonSink
        elif request.tap == 'postage':
            datasink = PostageCaptureSink
        else:
            raise ValueError(f'Unknown tap location {request.tap}')

        self.datasink = datasink(request, start=False)

        if submit:
            self.submit()
        elif isinstance(submit, tuple):
            self.submit(**submit)

    def __del__(self):
        self._kill_workers(kill_status_monitor=True, kill_data_sink=True)

    def status(self):
        """ Return the last known status of the request """
        return self._status_listener.latest()

    def status_history(self) -> tuple[str]:
        """Return a tuple of the status history of the job"""
        return self._status_listener.history()

    def cancel(self, kill_status_monitor=False, kill_data_sink=True):
        """

        Args:
            kill_status_monitor: Whether to terminate the status monitor
            kill_data_sink: Whether to terminate the data saver
        Returns: The parsed json submission result

        """
        self._kill_workers(kill_status_monitor=kill_status_monitor, kill_data_sink=kill_data_sink)
        ctx = zmq.Context().instance()
        with ctx.socket(zmq.REQ) as s:
            s.connect(self.request.server.command_url)
            s.send_pyobj(('abort', self.request.id))
            return s.recv_pyobj()

    def _kill_workers(self, kill_status_monitor=True, kill_data_sink=True):
        if kill_status_monitor and self._status_listener.is_alive():
            try:
                self.datasink.kill()
            except zmq.ZMQError as e:
                getLogger(__name__).warning(f'Caught {e} when telling data sink to terminate')
            if self._status_listener.is_alive():
                getLogger(__name__).warning(f'Status listener did not instantly terminate')
        if kill_data_sink and self.datasink.is_alive():
            try:
                self.datasink.kill()
            except zmq.ZMQError as e:
                getLogger(__name__).warning(f'Caught {e} when telling data sink to terminate')
            if self.datasink.is_alive():
                getLogger(__name__).warning(f'Data sink did not instantly terminate')

    def data(self, timeout=None):
        return self.datasink.data(timeout=timeout)

    def submit(self, with_sink=False, with_status=True):
        if with_status:
            self._status_listener.start()
            self._status_listener.ready()

        if with_sink and not self.request.data_endpoint.startswith('file://'):
            self.datasink.start()
            self.datasink.ready()

        try:
            ret = self.request.submit()
            self.submitted = True
            return ret
        except Exception as e:
            self._status_listener.kill()
            self.datasink.kill()
            raise e
