# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Copyright 2021 Daniel Mark Gass, see __about__.py for license information.
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
"""Array to bytes and bytes to array transform."""

from typing import Any, Generator, List, Optional, Sequence, Tuple, Type, Union

from .data import Data, DataMeta
from .dump import Record
from .int import IntX
from .transform import Transform

Dims = Sequence[Optional[int]]
DimsArg = Union[Dims, "ArrayX", IntX]
Format = Union[Type[Data], Transform]

GREEDY_DIMS: Dims = (None,)


def iterate_dims(array: Sequence[Any], ndims: int) -> Generator[int, None, None]:
    """Determine dimensions from multi-dimensional array value."""
    for _ in range(ndims):
        yield len(array)
        array = array[0]


class ArrayX(Transform):

    """Array to bytes and bytes to array transform."""

    __ndims__: int
    __df__: Tuple[DimsArg, Format]

    def __init__(self, name: str, fmt: Format, dims: Optional[DimsArg] = None) -> None:
        super().__init__(name)

        assert isinstance(fmt, (DataMeta, Transform))

        locked_dims: DimsArg
        ndims = 0

        if dims is None:
            locked_dims = GREEDY_DIMS

        elif isinstance(dims, IntX):
            locked_dims = dims

        elif isinstance(dims, ArrayX):
            locked_dims = dims
            if not dims.__ndims__:
                raise TypeError(
                    "invalid ArrayX 'dims', must be single dim and non-greedy"
                )

        else:
            locked_dims = tuple(None if d is None else int(d) for d in dims)
            assert all(True if d is None else d > 0 for d in locked_dims)
            if (len(locked_dims) == 1) and locked_dims[0] is not None:
                ndims = locked_dims[0]

        nbytes: Optional[int] = fmt.__nbytes__

        if nbytes is not None:
            if isinstance(locked_dims, tuple):
                for dim in locked_dims:
                    if dim is None:
                        nbytes = None
                        break

                    nbytes *= dim

            else:
                nbytes = None

        self.__df__ = locked_dims, fmt
        self.__ndims__ = ndims
        self.__nbytes__ = nbytes

    @property
    def fmt(self) -> Format:
        """Array element format."""
        return self.__df__[1]

    @property
    def dims(self) -> DimsArg:
        """Array dimensions."""
        return self.__df__[0]

    def __unpack__(
        self,
        buffer: bytes,
        offset: int,
        dump: Optional[Record] = None,
        dims: Optional[Dims] = None,
    ) -> Tuple[List[Any], int]:
        # pylint: disable=arguments-differ,too-many-branches
        if dump is not None:
            return ArrayX.__unpack_and_dump__(self, buffer, offset, dump, dims)

        x_dims, fmt = self.__df__

        if dims is None:
            if isinstance(x_dims, tuple):
                dims = x_dims
                if None in dims[1:]:
                    raise TypeError(
                        "array unpack does not support greedy secondary dimensions"
                    )

            elif isinstance(x_dims, IntX):
                dim, offset = x_dims.__unpack__(buffer, offset, dump)
                dims = (dim,)

            elif isinstance(x_dims, ArrayX):
                dims, offset = x_dims.__unpack__(buffer, offset, dump)

            else:  # pragma: no cover
                raise RuntimeError(f"invalid 'dims' {x_dims!r}")

        array: List[Any] = []
        append = array.append

        this_dim, *item_dims = dims

        if this_dim is None:
            end = len(buffer)
            while offset < end:
                item, offset = fmt.__unpack__(buffer, offset, dump)
                append(item)

        elif item_dims:
            for _ in range(this_dim):
                item, offset = ArrayX.__unpack__(self, buffer, offset, dump, item_dims)
                append(item)
        else:
            for _ in range(this_dim):
                item, offset = fmt.__unpack__(buffer, offset, dump)
                append(item)

        return array, offset

    def __unpack_and_dump__(
        self,
        buffer: bytes,
        offset: int,
        dump: Record,
        dims: Optional[Dims] = None,
    ) -> Tuple[List[Any], int]:
        # pylint: disable=arguments-differ,too-many-locals,too-many-branches
        x_dims, fmt = self.__df__

        if dims is None:
            if isinstance(x_dims, tuple):
                dims = x_dims
                if None in dims[1:]:
                    raise TypeError(
                        "array unpack does not support greedy secondary dimensions"
                    )

            elif isinstance(x_dims, IntX):
                dim, offset = x_dims.__unpack__(
                    buffer, offset, dump.add_record(access="len()", fmt=x_dims)
                )
                dims = (dim,)

            elif isinstance(x_dims, ArrayX):
                dims, offset = x_dims.__unpack__(
                    buffer, offset, dump.add_record(access="--dims--", fmt=x_dims)
                )

            else:  # pragma: no cover
                raise RuntimeError(f"invalid 'dims' {x_dims!r}")

        array: List[Any] = []
        append = array.append

        this_dim, *item_dims = dims

        if this_dim is None:
            end = len(buffer)
            i = 0
            fmt_name = fmt.__name__
            while offset < end:
                item, offset = fmt.__unpack__(
                    buffer, offset, dump.add_record(access=f"[{i}]", fmt=fmt_name)
                )
                append(item)
                i += 1

        elif item_dims:
            for i in range(this_dim):
                item, offset = ArrayX.__unpack_and_dump__(
                    self, buffer, offset, dump.add_record(access=f"[{i}]"), item_dims
                )
                append(item)
        else:
            fmt_name = fmt.__name__
            for i in range(this_dim):
                item, offset = fmt.__unpack__(
                    buffer, offset, dump.add_record(access=f"[{i}]", fmt=fmt_name)
                )
                append(item)

        if not array:
            dump.value = []

        return array, offset

    def __pack__(
        self,
        value: Sequence[Any],
        pieces: List[bytes],
        dump: Optional[Record] = None,
        dims: Optional[Dims] = None,
    ) -> None:
        # pylint: disable=arguments-differ,too-many-branches
        if dump is not None:
            ArrayX.__pack_and_dump__(self, value, pieces, dump, dims)

        else:
            x_dims, fmt = self.__df__

            if dims is None:
                if isinstance(x_dims, tuple):
                    this_dim, *item_dims = x_dims

                elif isinstance(x_dims, IntX):
                    this_dim, item_dims = len(value), []
                    x_dims.__pack__(this_dim, pieces, dump)

                elif isinstance(x_dims, ArrayX):
                    dims = tuple(iterate_dims(value, x_dims.__ndims__))
                    this_dim, *item_dims = dims
                    x_dims.__pack__(dims, pieces, dump)

                else:  # pragma: no cover
                    raise RuntimeError(f"invalid 'dims' {x_dims!r}")

            else:
                this_dim, *item_dims = dims

            try:
                actual_length = len(value)
            except TypeError:
                raise TypeError("invalid array value, retrying with dump") from None

            if this_dim is not None and actual_length != this_dim:
                raise TypeError("invalid array value, retrying with dump")

            if item_dims:
                for item in value:
                    ArrayX.__pack__(self, item, pieces, dump, item_dims)
            else:
                for item in value:
                    fmt.__pack__(item, pieces, dump)

    def __pack_and_dump__(
        self,
        value: Sequence[Any],
        pieces: List[bytes],
        dump: Record,
        dims: Optional[Dims] = None,
    ) -> None:
        # pylint: disable=arguments-differ,too-many-branches
        x_dims, fmt = self.__df__

        if dims is None:
            if isinstance(x_dims, tuple):
                this_dim, *item_dims = x_dims

            elif isinstance(x_dims, IntX):
                this_dim, item_dims = len(value), []
                x_dims.__pack__(
                    this_dim, pieces, dump.add_record(access="len()", fmt=x_dims)
                )

            elif isinstance(x_dims, ArrayX):
                dims = tuple(iterate_dims(value, x_dims.__ndims__))
                this_dim, *item_dims = dims
                x_dims.__pack__(
                    dims, pieces, dump.add_record(access="--dims--", fmt=x_dims)
                )

            else:  # pragma: no cover
                raise RuntimeError(f"invalid 'dims' {x_dims!r}")

        else:
            this_dim, *item_dims = dims

        try:
            actual_length = len(value)
        except TypeError:
            dump.value = value
            raise TypeError(
                f"invalid value, expected iterable of "
                f'{"any " if this_dim is None else ""}'
                f'length{"" if this_dim is None else " " + str(this_dim)}'
                f", got non-iterable"
            ) from None

        if not actual_length:
            dump.value = []

        if this_dim is None:
            this_dim = actual_length

        if item_dims:
            for i, item in zip(range(this_dim), value):
                ArrayX.__pack_and_dump__(
                    self, item, pieces, dump.add_record(access=f"[{i}]"), item_dims
                )
        else:
            fmt_name = fmt.__name__

            for i, item in zip(range(this_dim), value):
                fmt.__pack__(
                    item, pieces, dump.add_record(access=f"[{i}]", fmt=fmt_name)
                )

        if actual_length != this_dim:
            for i in range(actual_length, this_dim):
                dump.add_record(access=f"[{i}]", value="<missing>")

            for i, item in zip(range(this_dim, actual_length), value[this_dim:]):
                dump.add_record(
                    access=f"[{i}] <extra>", value=item, separate=(i == this_dim)
                )
            raise TypeError(
                f"invalid value, expected iterable of "
                f"{this_dim} length, got iterable of length {actual_length}"
            )
