import math
import struct
from functools import wraps
from contextlib import contextmanager
from typing import Any, Callable, Sequence
from collections.abc import Iterable
from . import ValueOutOfRangeException


def _usgn(i, width):
    """ Return i as an unsigned of given width, raising exceptions for out of
        bounds """
    i = int(i)
    if 0 <= i < 2**width:
        return i

    raise ValueOutOfRangeException(
        "%d doesn't fit in %d unsigned bits" % (i, width))


def _sgn(i, width):
    """ Return the unsigned that, when interpretted with given width,
        represents the signed value i """
    i = int(i)
    if i < -2**(width - 1) or 2**(width - 1) - 1 < i:
        raise ValueOutOfRangeException(
            "%d doesn't fit in %d signed bits" % (i, width))

    if i >= 0:
        return i

    return 2**width + i


def _upsgn(i, width):
    """ Return the signed integer that comes about by interpretting *i*
    as a signed field of *width* bytes"""
    i = int(i)
    if i < 0 or i > 2**width:
        raise ValueOutOfRangeException()

    if i < 2**(width - 1):
        return i

    return i - 2**width


class _reg_prop(object):
    def __init__(self, getter, setter, regs, length, offset, mask, _type, readonly=False):
        self.getter = getter
        self.setter = setter
        self.regs = regs
        self.length = length
        self.offset = offset
        self.mask = mask
        self._type = _type
        self.readonly = readonly

    def dump(self, inst, offset=0):
        return ([r + offset for r in self.regs], self.length, self.offset,
                self.mask, self._type, self.readonly, self.getter(inst))


def _readregs(inst, regs):
    """ read multiple regs, combine into one int
        LSBs first in regs
    """
    r = 0
    for idx in reversed(regs):
        r = (r << 32) | inst.readreg(idx)
    return r


def _writeregs(inst, regs, val):
    """ write one value spanning multiple regs
        LSBs first in regs
    """
    for idx in regs:
        inst.writereg(idx, val & 0xFFFFFFFF)
        val = val >> 32


def reg_sgnd(reg, offset=0, length=32, readonly=False):
    try:
        # check if reg is a list, otherise make it one
        iter(reg)
    except TypeError:
        reg = [reg]

    mask = ((1 << length) - 1) << offset

    def getter(s):
        return _upsgn((_readregs(s, reg) & mask) >> offset, length)

    def setter(s, v):
        r = _readregs(s, reg)
        v = _sgn(v, length) << offset
        v = (r & ~mask) | v
        return _writeregs(s, reg, v)

    return _reg_prop(getter, setter, reg, length, offset, mask, 'signed', readonly)


def reg_unsg(reg, offset=0, length=32, readonly=False):
    try:
        # check if reg is a list, otherise make it one
        iter(reg)
    except TypeError:
        reg = [reg]

    mask = ((1 << length) - 1) << offset

    def getter(s):
        r = _readregs(s, reg)
        return (r & mask) >> offset

    def setter(s, v):
        r = _readregs(s, reg)
        v = _usgn(v, length) << offset
        v = (r & ~mask) | v
        return _writeregs(s, reg, v)

    return _reg_prop(getter, setter, reg, length, offset, mask, 'unsigned', readonly)


def reg_slv(reg, offset=0, length=32, readonly=False):
    r = reg_unsg(reg, offset, length)
    r._type = 'slv'
    return r


def reg_float32(reg, readonly=False):
    def getter(s):
        return struct.unpack('<f', struct.pack('<I', s.readreg(reg)))[0]

    def setter(s, v):
        s.writereg(reg, struct.unpack('<I', struct.pack('<f', float(v)))[0])

    return _reg_prop(getter, setter, [reg], 32, 0, 0xFFFFFFFF, 'float', readonly)


def reg_bool(reg, offset, readonly=False):
    r = reg_unsg(reg, offset, 1, readonly)
    r._type = 'boolean'
    r.getter = lambda s, f=r.getter: bool(f(s))
    r.setter = lambda s, v, f=r.setter: f(s, bool(v))
    return r


def reg_enum(reg, enum_type, offset, length, readonly=False):
    mask = ((1 << length) - 1) << offset

    def getter(s):
        r = s.readreg(reg)
        v = (r & mask) >> offset
        return enum_type(v)

    def setter(s, v):
        try:
            v = enum_type(v)
        except ValueError:
            v = enum_type[v.upper()]

        r = s.readreg(reg)
        v = _usgn(v, length) << offset
        v = (r & ~mask) | v
        return s.writereg(reg, v)

    return _reg_prop(getter, setter, [reg], length, offset, mask, 'choice', readonly)


class _reg_prop_list(_reg_prop):
    """
    Represents an array of registers of any type. See RegisterAccess.reg_prop_list for details.
    """

    def __init__(
            self,
            parent: 'RegisterAccess',
            reg: int,
            offset: int,
            count: int,
            mkreg: Callable[..., _reg_prop],
            readonly: bool = False,
            padding: int = 0,
            **kwargs: Any):
        self.count = count
        self.offset = offset
        self.readonly = readonly
        self._parent = parent
        self.padding = padding

        assert count > 0
        assert 0 <= offset < 32, f'offset={offset} must be in [0, 32)'

        # Make the first item and use it to determine size, _type, etc.
        initial = mkreg(reg=reg, offset=offset, readonly=readonly, **kwargs)
        # The size actually taken up by the register
        item_size: int = initial.length
        # How much space is allocated for the register (which might overlap into the following word)
        size = item_size + padding

        assert size <= 32, f'_reg_prop_list: Cannot support items > 32 bits (size={size})'

        self.length = count * size
        self._type = 'list[' + initial._type + ']'
        # TODO: Does this make sense? It's not actually correct for padding != 0
        self.mask = ((1 << self.length) - 1) << offset
        num_regs = math.ceil(float(offset + self.length) / 32)
        self.regs = list(range(reg, reg + num_regs))

        self._reg_items: list[_reg_prop] = [initial]

        for i in range(1, count):
            bit_start = offset + i * size
            reg_offset, bit_offset = divmod(bit_start, 32)

            assert bit_offset + item_size <= 32, f'list elements must not straddle 32 bit '\
                f'boundaries, element {i} overlaps end of register (register={reg_offset}, '\
                f'offset={bit_offset}, size={size}, item_size item_size={item_size})'

            self._reg_items.append(mkreg(reg=reg + reg_offset,
                                         offset=bit_offset,
                                         readonly=readonly,
                                         **kwargs))

    def getter(self, ra: 'RegisterAccess'):
        return self

    def setter(self, ra: 'RegisterAccess', value: Sequence[Any]):
        assert len(value) == len(self._reg_items), \
            '_reg_prop_list.setter: List must have the same length'

        for i, v in enumerate(value):
            self._reg_items[i].setter(ra, v)

    def __getitem__(self, index) -> Any:
        res = self._reg_items[index]
        if isinstance(res, Iterable):
            return [r.getter(self._parent) for r in res]
        else:
            return res.getter(self._parent)

    def __setitem__(self, index, value):
        if isinstance(index, int):
            self._reg_items[index].setter(self._parent, value)
        else:
            for i, v in zip(range(index.start or 0,
                                  index.stop or len(self._reg_items),
                                  index.step or 1),
                            value):
                self._reg_items[i].setter(self._parent, v)

    def __len__(self) -> int:
        return len(self._reg_items)

    def __iter__(self):
        for r in self._reg_items:
            yield r.getter(self._parent)

    def dump(self, inst, offset=0):
        return ([r + offset for r in self.regs], self.length, self.offset,
                self.mask, self._type, self.readonly, [r.getter(inst) for r in self._reg_items])


class _reg_props(_reg_prop_list):
    def __init__(self, parent, regs):
        self._parent = parent
        self._reg_items = regs

    def dump(self, inst, offset=0):
        return [r.dump(inst, offset) for r in self._reg_items]


class RegisterAccess(object):
    def __init__(self, parent=None, offset=0, prefix=None):
        self._children = {}
        self.__parent = parent
        self.__offset = offset
        self.__prefix = prefix or type(self).__name__

        if self.__parent is not None:
            self.__parent.add_child(self, self.__prefix)

    def reg_control(self):
        if self.__parent is not None:
            return self.__parent.reg_control()
        return self

    def readreg(self, reg):
        if self.__parent is not None:
            return self.__parent.readreg(self.__offset + reg)
        raise NotImplementedError()

    def readblock(self, addr, length):
        if self.__parent is not None:
            # blocks are always absolute
            return self.__parent.readblock(addr, length)
        raise NotImplementedError()

    def writereg(self, reg, val):
        if self.__parent is not None:
            return self.__parent.writereg(self.__offset + reg, val)
        raise NotImplementedError()

    def writeblock(self, addr, data):
        if self.__parent is not None:
            # blocks are always absolute
            return self.__parent.writeblock(addr, data)
        raise NotImplementedError()

    def __getattribute__(self, name):
        attr = object.__getattribute__(self, name)
        if isinstance(attr, _reg_prop):
            return attr.getter(self)
        return attr

    def __setattr__(self, name, value):
        attr = None

        try:
            attr = object.__getattribute__(self, name)
        except AttributeError:
            pass

        if isinstance(attr, _reg_prop):
            assert not isinstance(value, _reg_prop)
            return attr.setter(self, value)

        super().__setattr__(name, value)

    def get_reg_prop(self, name):
        """ return the reg_prop instance without evaluating it
        """
        return object.__getattribute__(self, name)

    def reg_sgnd(self, reg, offset=0, length=32, readonly=False):
        return reg_sgnd(reg, offset, length, readonly)

    def reg_unsg(self, reg, offset=0, length=32, readonly=False):
        return reg_unsg(reg, offset, length, readonly)

    def reg_slv(self, reg, offset=0, length=32, readonly=False):
        return reg_slv(reg, offset, length, readonly)

    def reg_float32(self, reg, readonly=False):
        return reg_float32(reg, readonly)

    def reg_bool(self, reg, offset, readonly=False):
        return reg_bool(reg, offset, readonly)

    def reg_enum(self, reg, enum_type, offset, length, readonly=False):
        return reg_enum(reg, enum_type, offset, length, readonly)

    def reg_prop_list(
            self,
            reg: int,
            offset: int,
            count: int,
            mkreg: Callable[..., _reg_prop],
            readonly: bool = False,
            padding: int = 0,
            **kwargs):
        """
        Creates an array of register values of the type returned by mkreg.

        Keyword arguments:
        :param reg: the starting register
        :param offset: the bit offset into the starting register
        :param count: the number of registers to allocate
        :param mkreg: allocation function, see below
        :param readonly: passed to mkreg, whether these are read only values
        :param padding: padding in bits, to place after each register; see examples below. (default: 0)
        :param kwargs: passed to mkreg to configure it
        :returns: an object which contains all the registers, which can be accessed using list indexing/slicing.

        mkreg must be a function which returns a _reg_prop, generally one of self.reg_foo, and accepts arguments:
        * reg: int
        * offset: int
        * readonly: bool (default: False)
        * kwargs: any arguments specific to the register constructor.

        Examples:

        Allocate 5 8bit registers, containing ProtocolType enums, placed in bits 16-23 of the register
        leaving 24 bits of padding (this will use 5 32bit words, with the lower half and top 8 bits unallocated)

            |--------%%%%%%%%----------------||--------%%%%%%%%----------------|...
            31                               031                               0...

            def __init__(self,...):

                self.protos = self.reg_prop_list(4, 16, 5, self.reg_enum, padding=24, enum_type=ProtocolType, length=8)

            def set_coeff(self, i: int, val: ProtocolType):
                self.protos[i] = val

            def get_protos(self) -> Sequence[ProtocolType]:
                return self.protos

        This can be used to define interleving registers

            |--------%%%%%%%%$$$$$$$$$$$$$$$$||--------%%%%%%%%$$$$$$$$$$$$$$$$|...
            31                               031                               0...

        %: ProtocolType, $: coeffs, -: unallocated

            def __init__(self,...):
                ...
                self.coeffs = self.reg_prop_list(4, 0, 5, self.reg_sgnd, padding=16, length=16)

        """

        return _reg_prop_list(self, reg, offset, count, mkreg, readonly, padding, **kwargs)

    def reg_props(self, regs):
        return _reg_props(self, regs)

    def add_child(self, child, prefix):
        assert isinstance(child, RegisterAccess)
        self._children[f'{prefix}'] = child

    def dump_regs(self, offset=0):
        d = {}
        for k in dir(self):
            attr = object.__getattribute__(self, k)
            if isinstance(attr, _reg_prop):
                d[k] = attr.dump(self, offset=offset + self.__offset)

        for prefix, child in self._children.items():
            child_regs = child.dump_regs(offset=offset + self.__offset)
            d.update([(f'{prefix}.{k}', v) for k, v in child_regs.items()])

        return d

    def set_reg_by_name(self, path, val):
        curr = path.pop(0)
        if curr in self._children:
            return self._children[curr].set_reg_by_name(path, val)

        attr = object.__getattribute__(self, curr)
        if isinstance(attr, _reg_prop):
            return attr.setter(self, val)


class RegisterAccessControl(RegisterAccess):
    # This is handy for debugging unwanted register writes
    _assert_writes = False

    def __init__(self, region, control):
        super().__init__()
        self._region = region
        self._control = control

    def writereg(self, reg, val):
        assert not self._assert_writes
        self._control.writereg(self._region, reg, val)

    def writeregs(self, regs):
        # regs = [(0, 123), (1, 456), (addr, val), ...]
        assert not self._assert_writes
        self._control.writeregs(self._region, regs)

    def writeblock(self, addr, data):
        self._control.register_access([(self._region, True, addr, data)])

    def readreg(self, reg):
        return self._control.readreg(self._region, reg)

    def readregs(self, regs):
        # regs = [0, 1, 2, addr, ...]
        return self._control.readregs(self._region, regs)

    def readblock(self, addr, length):
        resp = self._control.register_access([(self._region, False, addr, length)])
        return resp[0][3]


def cached(f):
    """ Decorating a RegisterAccess method with @cached will cause
        any call to that method to first fetch the current state
        of the register map, and make all read/write interactions
        to the local cache.  Once the method exits, all write
        commands are flushed in a single packet, and in the order
        they were made
    """
    @wraps(f)
    def func(*args, **kwargs):
        inst = args[0]
        # propgate up the hierachy
        while not isinstance(inst, CachableAccessControl):
            inst = inst._RegisterAccess__parent

        with inst.cached_registers():
            return f(*args, **kwargs)
    return func


class CachableAccessControl(RegisterAccessControl):
    def __init__(self, region, control):
        super().__init__(region, control)
        self.__cached = False
        self.__num_regs = 128  # HACK

    def _fetch_state(self):
        block = self.readblock(0, 4 * self.__num_regs)
        self.__cached_regs = list(struct.unpack(f'<{self.__num_regs:d}I', block))
        self.__to_write = []

    def _close_state(self):
        if self.__to_write:
            super().writeregs(self.__to_write)

    def readreg(self, reg):
        if self.__cached:
            return self.__cached_regs[reg]
        return super().readreg(reg)

    def readregs(self, regs):
        if self.__cached:
            return [(r, self.readreg(r)) for r in regs]
        return super().readregs(regs)

    def writereg(self, reg, val):
        assert not self._assert_writes
        if self.__cached:
            self.__cached_regs[reg] = val
            # write ordering is important so we update the local cache
            # and append the command to write later
            self.__to_write.append((reg, val))
        else:
            return super().writereg(reg, val)

    def writeregs(self, regs):
        if self.__cached:
            for reg, val in regs:
                self.writereg(reg, val)
        else:
            super().writeregs(regs)

    @contextmanager
    def cached_registers(self, write=True):
        c = self.__cached  # local copy to maintain reentrant scope
        if not c:
            self._fetch_state()
            self.__cached = True

        try:
            yield
        finally:
            # don't write if an error occured
            if not c:
                self.__cached = False

        if write and not c:
            self._close_state()

    @cached
    def dump_regs(self):
        return super().dump_regs()


class RegisterAccessDummy(RegisterAccess):
    def __init__(self, length=0x1000):
        super().__init__()
        self.data = bytearray(length)

    def writereg(self, reg, val):
        struct.pack_into('<I', self.data, reg * 4, val)

    def writeregs(self, regs):
        for r, v in regs:
            self.writereg(r, v)

    def writeblock(self, addr, data):
        self.data[addr:addr + len(data)] = data

    def readreg(self, reg):
        return struct.unpack_from('<I', self.data, offset=reg * 4)[0]

    def readregs(self, regs):
        return [self.readreg(r) for r in regs]

    def readblock(self, addr, length):
        return self.data[addr:addr + length]
