from __future__ import annotations

import functools
from typing import Optional, Dict, Any, Union, Tuple, Callable
import inspect
import bitstring
from bitstring import utils

CACHE_SIZE = 256


def scaled_get_fn(get_fn, s: Union[int, float]):
    def wrapper(*args, scale=s, **kwargs):
        return get_fn(*args, **kwargs) * scale
    return wrapper


def scaled_set_fn(set_fn, s: Union[int, float]):
    def wrapper(bs, value, *args, scale=s, **kwargs):
        return set_fn(bs, value / scale, *args, **kwargs)
    return wrapper


def scaled_read_fn(read_fn, s: Union[int, float]):
    def wrapper(*args, scale=s, **kwargs):
        val = read_fn(*args, **kwargs)
        if isinstance(val, tuple):
            val, pos = val
            return val * scale, pos
        return val * scale
    return wrapper


class Dtype:
    """A data type class, representing a concrete interpretation of binary data.

    Dtype instances are immutable. They are often created implicitly elsewhere via a token string.

    >>> u12 = Dtype('uint', 12)  # length separate from token string.
    >>> float16 = Dtype('float16')  # length part of token string.
    >>> mxfp = Dtype('e3m2mxfp', scale=2 ** 6)  # dtype with scaling factor

    """

    _name: str
    _read_fn: Callable
    _set_fn: Callable
    _get_fn: Callable
    _return_type: Any
    _is_signed: bool
    _set_fn_needs_length: bool
    _variable_length: bool
    _bitlength: Optional[int]
    _bits_per_item: int
    _length: Optional[int]
    _scale: Union[None, float, int]

    def __new__(cls, token: Union[str, Dtype], /, length: Optional[int] = None, scale: Union[None, float, int] = None) -> Dtype:
        if isinstance(token, cls):
            return token
        if length is None:
            x = cls._new_from_token(token, scale)
            return x
        else:
            x = dtype_register.get_dtype(token, length, scale)
            return x

    @property
    def scale(self) -> Union[int, float, None]:
        """The multiplicative scale applied when interpreting the data."""
        return self._scale

    @property
    def name(self) -> str:
        """A string giving the name of the data type."""
        return self._name

    @property
    def length(self) -> int:
        """The length of the data type in units of bits_per_item. Set to None for variable length dtypes."""
        return self._length

    @property
    def bitlength(self) -> Optional[int]:
        """The number of bits needed to represent a single instance of the data type. Set to None for variable length dtypes."""
        return self._bitlength

    @property
    def bits_per_item(self) -> int:
        """The number of bits for each unit of length. Usually 1, but equals 8 for bytes type."""
        return self._bits_per_item

    @property
    def variable_length(self) -> bool:
        """If True then the length of the data type depends on the data being interpreted, and must not be specified."""
        return self._variable_length

    @property
    def return_type(self) -> Any:
        """The type of the value returned by the parse method, such as int, float or str."""
        return self._return_type

    @property
    def is_signed(self) -> bool:
        """If True then the data type represents a signed quantity."""
        return self._is_signed

    @property
    def set_fn(self) -> Optional[Callable]:
        """A function to set the value of the data type."""
        return self._set_fn

    @property
    def get_fn(self) -> Callable:
        """A function to get the value of the data type."""
        return self._get_fn

    @property
    def read_fn(self) -> Callable:
        """A function to read the value of the data type."""
        return self._read_fn

    def _set_scale(self, value: Union[None, float, int]) -> None:
        self._scale = value
        if self._scale is None:
            return
        if self._scale == 0:
            raise ValueError("A Dtype's scale factor must not be zero.")
        if not hasattr(self, 'unscaled_get_fn'):
            self.unscaled_get_fn = self._get_fn
            self.unscaled_set_fn = self._set_fn
            self.unscaled_read_fn = self._read_fn
        self._get_fn = scaled_get_fn(self.unscaled_get_fn, self._scale)
        self._set_fn = scaled_set_fn(self.unscaled_set_fn, self._scale)
        self._read_fn = scaled_read_fn(self.unscaled_read_fn, self._scale)

    @classmethod
    @functools.lru_cache(CACHE_SIZE)
    def _new_from_token(cls, token: str, scale: Union[None, float, int] = None) -> Dtype:
        token = ''.join(token.split())
        return dtype_register.get_dtype(*utils.parse_name_length_token(token), scale=scale)

    def __hash__(self) -> int:
        return hash((self._name, self._length))

    @classmethod
    @functools.lru_cache(CACHE_SIZE)
    def _create(cls, definition: DtypeDefinition, length: Optional[int], scale: Union[None, float, int]) -> Dtype:
        x = super().__new__(cls)
        x._name = definition.name
        x._bitlength = x._length = length
        x._bits_per_item = definition.multiplier
        if x._bitlength is not None:
            x._bitlength *= x._bits_per_item
        x._set_fn_needs_length = definition.set_fn_needs_length
        x._variable_length = definition.variable_length
        if x._variable_length or dtype_register.names[x._name].allowed_lengths.only_one_value():
            x._read_fn = definition.read_fn
        else:
            x._read_fn = functools.partial(definition.read_fn, length=x._bitlength)
        if definition.set_fn is None:
            x._set_fn = None
        else:
            if x._set_fn_needs_length:
                x._set_fn = functools.partial(definition.set_fn, length=x._bitlength)
            else:
                x._set_fn = definition.set_fn
        x._get_fn = definition.get_fn
        x._return_type = definition.return_type
        x._is_signed = definition.is_signed
        x._set_scale(scale)
        return x

    def build(self, value: Any, /) -> bitstring.Bits:
        """Create a bitstring from a value.

        The value parameter should be of a type appropriate to the dtype.
        """
        b = bitstring.Bits()
        self._set_fn(b, value)
        if self.bitlength is not None and len(b) != self.bitlength:
            raise ValueError(f"Dtype has a length of {self.bitlength} bits, but value '{value}' has {len(b)} bits.")
        return b

    def parse(self, b: BitsType, /) -> Any:
        """Parse a bitstring to find its value.

        The b parameter should be a bitstring of the appropriate length, or an object that can be converted to a bitstring."""
        b = bitstring.Bits._create_from_bitstype(b)
        return self._get_fn(bitstring.Bits(b))

    def __str__(self) -> str:
        if self._scale is not None:
            return self.__repr__()
        hide_length = self._variable_length or dtype_register.names[self._name].allowed_lengths.only_one_value() or self._length is None
        length_str = '' if hide_length else str(self._length)
        return f"{self._name}{length_str}"

    def __repr__(self) -> str:
        hide_length = self._variable_length or dtype_register.names[self._name].allowed_lengths.only_one_value() or self._length is None
        length_str = '' if hide_length else ', ' + str(self._length)
        if self._scale is None:
            scale_str = ''
        else:
            try:
                # This will only succeed for powers of two from -127 to 127.
                e8m0 = bitstring.Bits(e8m0mxfp=self._scale)
            except ValueError:
                scale_str = f', scale={self._scale}'
            else:
                power_of_two = e8m0.uint - 127
                if power_of_two in [0, 1]:
                    scale_str = f', scale={self._scale}'
                else:
                    scale_str = f', scale=2 ** {power_of_two}'
        return f"{self.__class__.__name__}('{self._name}'{length_str}{scale_str})"

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, Dtype):
            return self._name == other._name and self._length == other._length
        return False


class AllowedLengths:
    def __init__(self, value: Tuple[int, ...] = tuple()) -> None:
        if len(value) >= 3 and value[-1] is Ellipsis:
            step = value[1] - value[0]
            for i in range(1, len(value) - 1):
                if value[i] - value[i - 1] != step:
                    raise ValueError(f"Allowed length tuples must be equally spaced when final element is Ellipsis, but got {value}.")
            self.values = (value[0], value[1], Ellipsis)
        else:
            self.values = value

    def __str__(self) -> str:
        if self.values and self.values[-1] is Ellipsis:
            return f"({self.values[0]}, {self.values[1]}, ...)"
        return str(self.values)

    def __contains__(self, other: Any) -> bool:
        if not self.values:
            return True
        if self.values[-1] is Ellipsis:
            return (other - self.values[0]) % (self.values[1] - self.values[0]) == 0
        return other in self.values

    def only_one_value(self) -> bool:
        return self.values and len(self.values) == 1


class DtypeDefinition:
    """Represents a class of dtypes, such as uint or float, rather than a concrete dtype such as uint8.
    Not (yet) part of the public interface."""

    def __init__(self, name: str, set_fn, get_fn, return_type: Any = Any, is_signed: bool = False, bitlength2chars_fn=None,
                 variable_length: bool = False, allowed_lengths: Tuple[int, ...] = tuple(), multiplier: int = 1, description: str = ''):

        # Consistency checks
        if int(multiplier) != multiplier or multiplier <= 0:
            raise ValueError("multiplier must be an positive integer")
        if variable_length and allowed_lengths:
            raise ValueError("A variable length dtype can't have allowed lengths.")
        if variable_length and set_fn is not None and 'length' in inspect.signature(set_fn).parameters:
            raise ValueError("A variable length dtype can't have a set_fn which takes a length.")

        self.name = name
        self.description = description
        self.return_type = return_type
        self.is_signed = is_signed
        self.variable_length = variable_length
        self.allowed_lengths = AllowedLengths(allowed_lengths)

        self.multiplier = multiplier

        # Can work out if set_fn needs length based on its signature.
        self.set_fn_needs_length = set_fn is not None and 'length' in inspect.signature(set_fn).parameters
        self.set_fn = set_fn

        if self.allowed_lengths.values:
            def allowed_length_checked_get_fn(bs):
                if len(bs) not in self.allowed_lengths:
                    if self.allowed_lengths.only_one_value():
                        raise bitstring.InterpretError(f"'{self.name}' dtypes must have a length of {self.allowed_lengths.values[0]}, but received a length of {len(bs)}.")
                    else:
                        raise bitstring.InterpretError(f"'{self.name}' dtypes must have a length in {self.allowed_lengths}, but received a length of {len(bs)}.")
                return get_fn(bs)
            self.get_fn = allowed_length_checked_get_fn  # Interpret everything and check the length
        else:
            self.get_fn = get_fn  # Interpret everything

        # Create a reading function from the get_fn.
        if not self.variable_length:
            if self.allowed_lengths.only_one_value():
                def read_fn(bs, start):
                    return self.get_fn(bs[start:start + self.allowed_lengths.values[0]])
            else:
                def read_fn(bs, start, length):
                    if len(bs) < start + length:
                        raise bitstring.ReadError(f"Needed a length of at least {length} bits, but only {len(bs) - start} bits were available.")
                    return self.get_fn(bs[start:start + length])
            self.read_fn = read_fn
        else:
            # We only find out the length when we read/get.
            def length_checked_get_fn(bs):
                x, length = get_fn(bs)
                if length != len(bs):
                    raise ValueError
                return x
            self.get_fn = length_checked_get_fn

            def read_fn(bs, start):
                try:
                    x, length = get_fn(bs[start:])
                except bitstring.InterpretError:
                    raise bitstring.ReadError
                return x, start + length
            self.read_fn = read_fn
        self.bitlength2chars_fn = bitlength2chars_fn

    def get_dtype(self, length: Optional[int] = None, scale: Union[None, float, int] = None) -> Dtype:
        if self.allowed_lengths:
            if length is None:
                if self.allowed_lengths.only_one_value():
                    length = self.allowed_lengths.values[0]
            else:
                if length not in self.allowed_lengths:
                    if self.allowed_lengths.only_one_value():
                        raise ValueError(f"A length of {length} was supplied for the '{self.name}' dtype, but its only allowed length is {self.allowed_lengths.values[0]}.")
                    else:
                        raise ValueError(f"A length of {length} was supplied for the '{self.name}' dtype which is not one of its possible lengths (must be one of {self.allowed_lengths}).")
        if length is None:
            d = Dtype._create(self, None, scale)
            return d
        if self.variable_length:
            raise ValueError(f"A length ({length}) shouldn't be supplied for the variable length dtype '{self.name}'.")
        d = Dtype._create(self, length, scale)
        return d

    def __repr__(self) -> str:
        s = f"{self.__class__.__name__}(name='{self.name}', description='{self.description}', return_type={self.return_type.__name__}, "
        s += f"is_signed={self.is_signed}, set_fn_needs_length={self.set_fn_needs_length}, allowed_lengths={self.allowed_lengths!s}, multiplier={self.multiplier})"
        return s


class Register:
    """A singleton class that holds all the DtypeDefinitions. Not (yet) part of the public interface."""

    _instance: Optional[Register] = None
    names: Dict[str, DtypeDefinition] = {}

    def __new__(cls) -> Register:
        # Singleton. Only one Register instance can ever exist.
        if cls._instance is None:
            cls._instance = super(Register, cls).__new__(cls)
        return cls._instance

    @classmethod
    def add_dtype(cls, definition: DtypeDefinition):
        cls.names[definition.name] = definition
        if definition.get_fn is not None:
            setattr(bitstring.bits.Bits, definition.name, property(fget=definition.get_fn, doc=f"The bitstring as {definition.description}. Read only."))
        if definition.set_fn is not None:
            setattr(bitstring.bitarray_.BitArray, definition.name, property(fget=definition.get_fn, fset=definition.set_fn, doc=f"The bitstring as {definition.description}. Read and write."))

    @classmethod
    def add_dtype_alias(cls, name: str, alias: str):
        cls.names[alias] = cls.names[name]
        definition = cls.names[alias]
        if definition.get_fn is not None:
            setattr(bitstring.bits.Bits, alias, property(fget=definition.get_fn, doc=f"An alias for '{name}'. Read only."))
        if definition.set_fn is not None:
            setattr(bitstring.bitarray_.BitArray, alias, property(fget=definition.get_fn, fset=definition.set_fn, doc=f"An alias for '{name}'. Read and write."))

    @classmethod
    def get_dtype(cls, name: str, length: Optional[int], scale: Union[None, float, int] = None) -> Dtype:
        try:
            definition = cls.names[name]
        except KeyError:
            raise ValueError(f"Unknown Dtype name '{name}'. Names available: {list(cls.names.keys())}.")
        else:
            return definition.get_dtype(length, scale)

    @classmethod
    def __getitem__(cls, name: str) -> DtypeDefinition:
        return cls.names[name]

    @classmethod
    def __delitem__(cls, name: str) -> None:
        del cls.names[name]

    def __repr__(self) -> str:
        s = [f"{'key':<12}:{'name':^12}{'signed':^8}{'set_fn_needs_length':^23}{'allowed_lengths':^16}{'multiplier':^12}{'return_type':<13}"]
        s.append('-' * 85)
        for key in self.names:
            m = self.names[key]
            allowed = '' if not m.allowed_lengths else m.allowed_lengths
            ret = 'None' if m.return_type is None else m.return_type.__name__
            s.append(f"{key:<12}:{m.name:>12}{m.is_signed:^8}{m.set_fn_needs_length:^16}{allowed!s:^16}{m.multiplier:^12}{ret:<13} # {m.description}")
        return '\n'.join(s)


# Create the Register singleton
dtype_register = Register()