import logging
from enum import IntEnum, Enum, auto

from pymoku import version, progress

from typing import NamedTuple, Optional

# hack to allow top level imports from pymoku namespace
from pkgutil import extend_path
__path__ = extend_path(__path__, __name__)

try:
    from itertools import batched  # only for python >= 3.12
except ImportError:
    from itertools import islice

    def batched(iterable, n, *, strict=False):
        # https://docs.python.org/3/library/itertools.html#itertools.batched
        # batched('ABCDEFG', 2) → AB CD EF G
        if n < 1:
            raise ValueError('n must be at least one')
        iterator = iter(iterable)
        while batch := tuple(islice(iterator, n)):
            if strict and len(batch) != n:
                raise ValueError('batched(): incomplete batch')
            yield batch


class dtype(Enum):
    s16x4 = auto()
    s32x2 = auto()
    s64x1 = auto()
    digital = auto()
    iq16x2 = auto()
    iq32x1 = auto()


class MokuInfo(NamedTuple):
    name: Optional[str] = None
    netver: Optional[int] = None
    fwver: Optional[int] = None
    hwver: Optional[float] = None
    serial: Optional[int] = None
    colour: Optional[str] = None
    bootmode: Optional[str] = None
    ip_addr: Optional[str] = None


__version__ = version.version

log = logging.getLogger('pymoku')


class MokuException(Exception):
    """
    Base class for other Exceptions
    """
    pass


class DeployException(MokuException):
    """
    Couldn't start instrument. Moku may not be licenced to use that instrument
    """
    pass


class ValueOutOfRangeException(MokuException):
    """
    Invalid value for this operation
    """
    pass


class StreamException(MokuException):
    def __init__(self, message, err=None):
        """
        Data logging was interrupted or failed
        """
        super(StreamException, self).__init__(message)
        self.err = err


class AInMode(IntEnum):
    DDS = 0x00
    PRECISION = 0x01
    MINMAX = 0X02
    ANDOR = 0X03
    PROTOCOL = 0X04


class LogProgressHandler(progress.ProgressHandler):
    def set_message(self, message):
        log.info(message)

    def complete(self):
        log.debug("Done")


progress.register_progress_handler(LogProgressHandler())


from pymoku.moku import Moku, MOKU_CLASS  # noqa
from pymoku.slot import Slot  # noqa
from pymoku.network import ZMQCurveSocketFactory  # noqa
from pymoku.network import Control, Control529, Control511, Control433, FrameTimeout  # noqa
from pymoku.network import FileTransfer, FileTransfer511, FileTransfer433  # noqa


def _try_connect(factory, fw_ver=None):
    """ Attempt to connect to the control socket
        fw_ver provides a hint on which network version to use
    """
    control = None

    if fw_ver is None or fw_ver > 529:
        try:
            control = Control(factory.get_control_socket())
            control.get_property('device.fw_version')
            return control
        except Exception:
            if control:
                control.close()

    if fw_ver is None or fw_ver > 511:
        try:
            log.debug('attempting to connect as 529')
            control = Control529(factory.get_control_socket())
            control.get_property('device.fw_version')
            return control
        except Exception:
            if control:
                control.close()

    if fw_ver is None or fw_ver > 433:
        try:
            log.debug('attempting to connect as 511')
            control = Control511(factory.get_control_socket())
            control.get_property('device.fw_version')
            return control
        except Exception:
            if control:
                control.close()
            raise

    if fw_ver is None:
        try:
            log.debug('attempting to connect as 433')
            control = Control433(factory.get_control_socket())
            control.get_property('device.fw_version')
            return control
        except Exception:
            if control:
                control.close()
            raise

    raise MokuException("Couldn't connect to Moku")


def _find(serial):
    from . import finder
    minfo = finder.Finder().find_all(timeout=1, filter=lambda minfo: minfo.serial == serial)
    if minfo:
        return minfo[0].ip_addr

    minfo = finder.Finder(ipv6=True).find_all(timeout=1, filter=lambda minfo: minfo.serial == serial)
    if minfo:
        return minfo[0].ip_addr

    raise MokuException("Can't find Moku with serial == {:d})".format(serial))


def connect(uid, socket_class=ZMQCurveSocketFactory, attach=True):
    """ Return a Moku instance suitable to the hardware version
    """
    fw_ver = None
    try:
        uid = int(uid)
    except (ValueError, TypeError):
        pass

    if isinstance(uid, MokuInfo):
        ip_addr = uid.ip_addr
        fw_ver = uid.fwver
    elif isinstance(uid, int):
        ip_addr = _find(uid)
    else:
        ip_addr = uid

    factory = socket_class(ip_addr)

    with progress.progress() as prog:
        prog.set_message('Connecting...')
        control = _try_connect(factory, fw_ver)

    fw_ver, hw_ver = control.get_properties(['device.fw_version', 'device.hw_version'])
    fw_ver = int(fw_ver['device.fw_version'])
    hw_ver = str(hw_ver['device.hw_version'])

    if fw_ver < 462:
        fileserv = FileTransfer433(control.skt)
    elif fw_ver < 512:
        fileserv = FileTransfer511(control.skt)
    else:
        fileserv = FileTransfer(control.skt)

    m = MOKU_CLASS[hw_ver](factory, control, fileserv)

    log.info(f"Connected to {type(m).__name__}[{m.serial:06d}]"
             f" \"{m.name}\" fw_ver: {m.fw_ver}")

    if m.bootmode == 'recovery':
        log.warning("Moku in recovery mode.  Run an update to fix")

    try:
        if attach:
            m.attach()
    except Exception:
        log.exception("Error initializing Moku")

    return m


def connect_tcp(uid, attach=True):
    from .network import TCPSocketFactory  # noqa
    return connect(uid, socket_class=TCPSocketFactory, attach=attach)


def connect_http(uid, attach=True):
    from .network import HTTPSocketFactory  # noqa
    return connect(uid, socket_class=HTTPSocketFactory, attach=attach)


def connect_ipc(uid=None, attach=True):
    from .network import ZMQIPCSocketFactory  # noqa
    return connect(uid, socket_class=ZMQIPCSocketFactory, attach=attach)
