from __future__ import annotations

import bitarray
import bitarray.util
from bitstring.exceptions import CreationError
from typing import Union, Iterable, Optional, overload, Iterator, Any
from bitstring.helpers import offset_slice_indices_lsb0

if bitarray.__version__.startswith("2."):
    raise ImportError(f"bitstring version 4.3 requires bitarray version 3 or higher. Found version {bitarray.__version__}.")


class _BitStore:
    """A light wrapper around bitarray that does the LSB0 stuff"""

    __slots__ = ('_bitarray', 'modified_length', 'immutable')

    def __init__(self, initializer: Union[bitarray.bitarray, None] = None,
                 immutable: bool = False) -> None:
        if isinstance(initializer, str):
            assert False
        self._bitarray = bitarray.bitarray(initializer)
        self.immutable = immutable
        self.modified_length = None

    @classmethod
    def from_zeros(cls, i: int) -> _BitStore:
        x = super().__new__(cls)
        x._bitarray = bitarray.bitarray(i)
        x.immutable = False
        x.modified_length = None
        return x


    @classmethod
    def from_bin(cls, s: str) -> _BitStore:
        x = super().__new__(cls)
        x._bitarray = bitarray.bitarray(s)
        x.immutable = False
        x.modified_length = None
        return x

    @classmethod
    def from_bytes(cls, b: Union[bytes, bytearray, memoryview], /) -> _BitStore:
        x = super().__new__(cls)
        x._bitarray = bitarray.bitarray()
        x._bitarray.frombytes(b)
        x.immutable = False
        x.modified_length = None
        return x

    @classmethod
    def frombuffer(cls, buffer, /, length: Optional[int] = None) -> _BitStore:
        x = super().__new__(cls)
        x._bitarray = bitarray.bitarray(buffer=buffer)
        x.immutable = True
        x.modified_length = length
        # Here 'modified' means it shouldn't be changed further, so setting, deleting etc. are disallowed.
        if x.modified_length is not None:
            if x.modified_length < 0:
                raise CreationError("Can't create bitstring with a negative length.")
            if x.modified_length > len(x._bitarray):
                raise CreationError(
                    f"Can't create bitstring with a length of {x.modified_length} from {len(x._bitarray)} bits of data.")
        return x

    @classmethod
    def join(cls, bitstores: Iterable[_BitStore], /) -> _BitStore:
        x = super().__new__(cls)
        x._bitarray = bitarray.bitarray()
        for b in bitstores:
            x._bitarray += b._bitarray
        x.immutable = False
        x.modified_length = None
        return x

    @staticmethod
    def using_rust_core() -> bool:
        return False

    def tobitarray(self) -> bitarray.bitarray:
        if self.modified_length is not None:
            return self.getslice(0, len(self))._bitarray
        return self._bitarray

    def to_bytes(self) -> bytes:
        if self.modified_length is not None:
            return self._bitarray[:self.modified_length].tobytes()
        return self._bitarray.tobytes()

    def to_u(self) -> int:
        if self.modified_length is not None:
            return bitarray.util.ba2int(self._bitarray[:self.modified_length], signed=False)
        return bitarray.util.ba2int(self._bitarray, signed=False)

    def to_i(self) -> int:
        if self.modified_length is not None:
            return bitarray.util.ba2int(self._bitarray[:self.modified_length], signed=True)
        return bitarray.util.ba2int(self._bitarray, signed=True)

    def to_hex(self) -> str:
        if self.modified_length is not None:
            return bitarray.util.ba2hex(self._bitarray[:self.modified_length])
        return bitarray.util.ba2hex(self._bitarray)

    def to_bin(self) -> str:
        if self.modified_length is not None:
            return self._bitarray[:self.modified_length].to01()
        return self._bitarray.to01()

    def to_oct(self) -> str:
        if self.modified_length is not None:
            return bitarray.util.ba2base(8, self._bitarray[:self.modified_length])
        return bitarray.util.ba2base(8, self._bitarray)

    def __imul__(self, n: int, /) -> _BitStore:
        self._bitarray *= n
        return self

    def __ilshift__(self, n: int, /) -> None:
        self._bitarray <<= n

    def __irshift__(self, n: int, /) -> None:
        self._bitarray >>= n

    def __iadd__(self, other: _BitStore, /) -> _BitStore:
        self._bitarray += other._bitarray
        return self

    def __add__(self, other: _BitStore, /) -> _BitStore:
        bs = self._mutable_copy()
        bs += other
        return bs

    def __eq__(self, other: Any, /) -> bool:
        return self._bitarray == other._bitarray

    def __and__(self, other: _BitStore, /) -> _BitStore:
        return _BitStore(self._bitarray & other._bitarray)

    def __or__(self, other: _BitStore, /) -> _BitStore:
        return _BitStore(self._bitarray | other._bitarray)

    def __xor__(self, other: _BitStore, /) -> _BitStore:
        return _BitStore(self._bitarray ^ other._bitarray)

    def __iand__(self, other: _BitStore, /) -> _BitStore:
        self._bitarray &= other._bitarray
        return self

    def __ior__(self, other: _BitStore, /) -> _BitStore:
        self._bitarray |= other._bitarray
        return self

    def __ixor__(self, other: _BitStore, /) -> _BitStore:
        self._bitarray ^= other._bitarray
        return self

    def __invert__(self) -> _BitStore:
        return _BitStore(~self._bitarray)

    def find(self, bs: _BitStore, start: int, end: int, bytealigned: bool = False) -> int | None:
        if not bytealigned:
            x = self._bitarray.find(bs._bitarray, start, end)
            return None if x == -1 else x
        try:
            return next(self.findall_msb0(bs, start, end, bytealigned))
        except StopIteration:
            return None

    def rfind(self, bs: _BitStore, start: int, end: int, bytealigned: bool = False) -> int | None:
        if not bytealigned:
            x = self._bitarray.find(bs._bitarray, start, end, right=True)
            return None if x == -1 else x
        try:
            return next(self.rfindall_msb0(bs, start, end, bytealigned))
        except StopIteration:
            return None

    def findall_msb0(self, bs: _BitStore, start: int, end: int, bytealigned: bool = False) -> Iterator[int]:
        if bytealigned is True and len(bs) % 8 == 0:
            # Special case, looking for whole bytes on whole byte boundaries
            bytes_ = bs.to_bytes()
            # Round up start byte to next byte, and round end byte down.
            # We're only looking for whole bytes, so can ignore bits at either end.
            start_byte = (start + 7) // 8
            end_byte = end // 8
            b = self._bitarray[start_byte * 8: end_byte * 8].tobytes()
            byte_pos = 0
            bytes_to_search = end_byte - start_byte
            while byte_pos < bytes_to_search:
                byte_pos = b.find(bytes_, byte_pos)
                if byte_pos == -1:
                    break
                yield (byte_pos + start_byte) * 8
                byte_pos = byte_pos + 1
            return
        # General case
        i = self._bitarray.search(bs._bitarray, start, end)
        if not bytealigned:
            for p in i:
                yield p
        else:
            for p in i:
                if (p % 8) == 0:
                    yield p

    def rfindall_msb0(self, bs: _BitStore, start: int, end: int, bytealigned: bool = False) -> Iterator[int]:
        i = self._bitarray.search(bs._bitarray, start, end, right=True)
        if not bytealigned:
            for p in i:
                yield p
        else:
            for p in i:
                if (p % 8) == 0:
                    yield p

    def count(self, value, /) -> int:
        return self._bitarray.count(value)

    def clear(self) -> None:
        self._bitarray.clear()

    def reverse(self) -> None:
        self._bitarray.reverse()

    def __iter__(self) -> Iterable[bool]:
        for i in range(len(self)):
            yield self.getindex(i)

    def _mutable_copy(self) -> _BitStore:
        """Always creates a copy, even if instance is immutable."""
        return _BitStore(self._bitarray, immutable=False)

    def as_immutable(self) -> _BitStore:
        return _BitStore(self._bitarray, immutable=True)

    def copy(self) -> _BitStore:
        return self if self.immutable else self._mutable_copy()

    def __getitem__(self, item: Union[int, slice], /) -> Union[int, _BitStore]:
        # Use getindex or getslice instead
        raise NotImplementedError

    def getindex_msb0(self, index: int, /) -> bool:
        return bool(self._bitarray.__getitem__(index))

    def getslice_withstep_msb0(self, key: slice, /) -> _BitStore:
        if self.modified_length is not None:
            key = slice(*key.indices(self.modified_length))
        return _BitStore(self._bitarray.__getitem__(key))

    def getslice_withstep_lsb0(self, key: slice, /) -> _BitStore:
        key = offset_slice_indices_lsb0(key, len(self))
        return _BitStore(self._bitarray.__getitem__(key))

    def getslice_msb0(self, start: Optional[int], stop: Optional[int], /) -> _BitStore:
        if self.modified_length is not None:
            key = slice(*slice(start, stop, None).indices(self.modified_length))
            start = key.start
            stop = key.stop
        return _BitStore(self._bitarray[start:stop])

    def getslice_lsb0(self, start: Optional[int], stop: Optional[int], /) -> _BitStore:
        s = offset_slice_indices_lsb0(slice(start, stop, None), len(self))
        return _BitStore(self._bitarray[s.start:s.stop])

    def getindex_lsb0(self, index: int, /) -> bool:
        return bool(self._bitarray.__getitem__(-index - 1))

    @overload
    def setitem_lsb0(self, key: int, value: int, /) -> None:
        ...

    @overload
    def setitem_lsb0(self, key: slice, value: _BitStore, /) -> None:
        ...

    def setitem_lsb0(self, key: Union[int, slice], value: Union[int, _BitStore], /) -> None:
        if isinstance(key, slice):
            new_slice = offset_slice_indices_lsb0(key, len(self))
            self._bitarray.__setitem__(new_slice, value._bitarray)
        else:
            self._bitarray.__setitem__(-key - 1, value)

    def delitem_lsb0(self, key: Union[int, slice], /) -> None:
        if isinstance(key, slice):
            new_slice = offset_slice_indices_lsb0(key, len(self))
            self._bitarray.__delitem__(new_slice)
        else:
            self._bitarray.__delitem__(-key - 1)

    def invert_msb0(self, index: Optional[int] = None, /) -> None:
        if index is not None:
            self._bitarray.invert(index)
        else:
            self._bitarray.invert()

    def invert_lsb0(self, index: Optional[int] = None, /) -> None:
        if index is not None:
            self._bitarray.invert(-index - 1)
        else:
            self._bitarray.invert()

    def extend_left(self, other: _BitStore, /) -> None:
        self._bitarray = other._bitarray + self._bitarray

    def any(self) -> bool:
        return self._bitarray.any()

    def all(self) -> bool:
        return self._bitarray.all()

    def __len__(self) -> int:
        return self.modified_length if self.modified_length is not None else len(self._bitarray)

    def setitem_msb0(self, key, value, /):
        if isinstance(value, _BitStore):
            self._bitarray.__setitem__(key, value._bitarray)
        else:
            self._bitarray.__setitem__(key, value)

    def delitem_msb0(self, key, /):
        self._bitarray.__delitem__(key)


ConstBitStore = _BitStore
MutableBitStore = _BitStore