from hub.constants import ENCODING_DTYPE, UUID_SHIFT_AMOUNT
from hub.util.exceptions import ChunkIdEncoderError
import hub
from hub.core.storage.cachable import Cachable
from io import BytesIO
from typing import Optional, Tuple
import numpy as np
from uuid import uuid4

# these constants are for accessing the data layout. see the `ChunkIdEncoder` docstring.
CHUNK_ID_INDEX = 0
LAST_INDEX_INDEX = 1


class ChunkIdEncoder(Cachable):
    def __init__(self):
        """Custom compressor that allows reading of chunk IDs from a sample index without decompressing.

        Chunk IDs:
            Chunk IDs are a `ENCODING_DTYPE` value  and this class handles generating/encoding them.

        Layout:
            `_encoded_ids` is a 2D array.

            Rows:
                The number of rows is equal to the number of chunk IDs this encoder is responsible for.

            Columns:
                The number of columns is 2.
                Each row looks like this: [chunk_id, last_index], where `last_index` is the last index that the
                chunk with `chunk_id` contains.

            Example:
                >>> enc = ChunkIdEncoder()
                >>> enc.generate_chunk_id()
                >>> enc.num_chunks
                1
                >>> enc.register_samples_to_last_chunk_id(10)
                >>> enc.num_samples
                10
                >>> enc.register_samples_to_last_chunk_id(10)
                >>> enc.num_samples
                20
                >>> enc.num_chunks
                1
                >>> enc.generate_chunk_id()
                >>> enc.register_samples_to_last_chunk_id(1)
                >>> enc.num_samples
                21
                >>> enc._encoded_ids
                [[3723322941, 19],
                 [1893450271, 20]]
                >>> enc[20]
                1893450271

            Best case scenario:
                The best case scenario is when all samples fit within a single chunk. This means the number of rows is 1,
                providing a O(1) lookup.

            Worst case scenario:
                The worst case scenario is when only 1 sample fits per chunk. This means the number of rows is equal to the number
                of samples, providing a O(log(N)) lookup.

            Lookup algorithm:
                To get the chunk ID for some sample index, you do a binary search over the right-most column. This will give you
                the row that corresponds to that sample index (since the right-most column is our "last index" for that chunk ID).
                Then, you get the left-most column and that is your chunk ID!

        """

        self._encoded_ids = None

    def tobytes(self) -> memoryview:
        bio = BytesIO()
        np.savez(
            bio,
            version=hub.__encoded_version__,
            ids=self._encoded_ids,
        )
        return bio.getbuffer()

    @staticmethod
    def name_from_id(id: ENCODING_DTYPE) -> str:
        """Returns the hex of `id` with the "0x" prefix removed. This is the chunk's name and should be used to determine the chunk's key.
        Can convert back into `id` using `id_from_name`. You can get the `id` for a chunk using `__getitem__`."""

        return hex(id)[2:]

    @staticmethod
    def id_from_name(name: str) -> ENCODING_DTYPE:
        """Returns the 64-bit integer from the hex `name` generated by `name_from_id`."""

        return int("0x" + name, 16)

    def get_name_for_chunk(self, chunk_index: int) -> str:
        """Gets the name for the chunk at index `chunk_index`. If you need to get the name for a chunk from a sample index, instead
        use `__getitem__`, then `name_from_id`."""

        chunk_id = self._encoded_ids[:, CHUNK_ID_INDEX][chunk_index]
        return ChunkIdEncoder.name_from_id(chunk_id)

    @classmethod
    def frombuffer(cls, buffer: bytes):
        instance = cls()
        bio = BytesIO(buffer)
        npz = np.load(bio)
        instance._encoded_ids = npz["ids"]
        return instance

    @property
    def num_chunks(self) -> int:
        if self._encoded_ids is None:
            return 0
        return len(self._encoded_ids)

    @property
    def num_samples(self) -> int:
        if self._encoded_ids is None:
            return 0
        return int(self._encoded_ids[-1, LAST_INDEX_INDEX] + 1)

    def generate_chunk_id(self) -> ENCODING_DTYPE:
        """Generates a random 64bit chunk ID using uuid4. Also prepares this ID to have samples registered to it.
        This method should be called once per chunk created.

        Returns:
            ENCODING_DTYPE: The random chunk ID.
        """

        id = ENCODING_DTYPE(uuid4().int >> UUID_SHIFT_AMOUNT)

        if self.num_samples == 0:
            self._encoded_ids = np.array([[id, -1]], dtype=ENCODING_DTYPE)

        else:
            last_index = self.num_samples - 1

            new_entry = np.array(
                [[id, last_index]],
                dtype=ENCODING_DTYPE,
            )
            self._encoded_ids = np.concatenate([self._encoded_ids, new_entry])

        return id

    def register_samples_to_last_chunk_id(self, num_samples: int):
        """Registers samples to the chunk ID that was generated last with the `generate_chunk_id` method.
        This method should be called at least once per chunk created.

        Args:
            num_samples (int): The number of samples the last chunk ID should have added to it's registration.

        Raises:
            ValueError: `num_samples` should be non-negative.
            ChunkIdEncoderError: Must call `generate_chunk_id` before registering samples.
            ChunkIdEncoderError: `num_samples` can only be 0 if it is able to be a sample continuation accross chunks.
        """

        if num_samples < 0:
            raise ValueError(
                f"Cannot register negative num samples. Got: {num_samples}"
            )

        if self.num_samples == 0:
            raise ChunkIdEncoderError(
                "Cannot register samples because no chunk IDs exist."
            )

        if num_samples == 0 and self.num_chunks < 2:
            raise ChunkIdEncoderError(
                "Cannot register 0 num_samples (signifying a partial sample continuing the last chunk) when no last chunk exists."
            )

        current_entry = self._encoded_ids[-1]

        # this operation will trigger an overflow for the first addition, so supress the warning
        np.seterr(over="ignore")
        current_entry[LAST_INDEX_INDEX] += ENCODING_DTYPE(num_samples)
        np.seterr(over="warn")

    def get_local_sample_index(self, global_sample_index: int) -> int:
        """Converts `global_sample_index` into a new index that is relative to the chunk the sample belongs to.

        Example:
            Given: 2 sampes in chunk 0, 2 samples in chunk 1, and 3 samples in chunk 2.

            >>> self.num_samples
            7
            >>> self.num_chunks
            3
            >>> self.get_local_sample_index(0)
            0
            >>> self.get_local_sample_index(1)
            1
            >>> self.get_local_sample_index(2)
            0
            >>> self.get_local_sample_index(3)
            1
            >>> self.get_local_sample_index(6)
            2

        Args:
            global_sample_index (int): Index of the sample relative to the containing tensor.

        Returns:
            int: local index value between 0 and the amount of samples the chunk contains - 1.
        """

        _, chunk_index = self.__getitem__(global_sample_index, return_chunk_index=True)  # type: ignore

        if chunk_index == 0:
            return global_sample_index

        current_entry = self._encoded_ids[chunk_index - 1]  # type: ignore
        last_num_samples = current_entry[LAST_INDEX_INDEX] + 1

        return int(global_sample_index - last_num_samples)

    def __getitem__(
        self, sample_index: int, return_chunk_index: bool = False
    ) -> Tuple[ENCODING_DTYPE, Optional[int]]:
        """Get the ID for the chunk that `sample_index` is stored in.
        To get the name of the chunk, use `name_from_id`.

        Args:
            sample_index (int): Global index (relative to the tensor). This will be converted to the local chunk index.
            return_chunk_index (bool): If True, 2 values are returned, the second one being the chunk's index. Defaults to False.

        Raises:
            IndexError: If no samples exist or `sample_index` exceeds the available indices.

        Returns:
            Tuple[Tuple[ENCODING_DTYPE], Optional[Tuple[int]]]: Returns the chunk ID for `sample_index`. If `return_chunk_index` is True,
                there will be 2 values. The second one being the chunk's index.
        """

        if self.num_samples == 0:
            raise IndexError(
                f"Index {sample_index} is out of bounds for an empty chunk names encoding."
            )

        if sample_index < 0:
            sample_index = (self.num_samples) + sample_index

        idx = np.searchsorted(self._encoded_ids[:, LAST_INDEX_INDEX], sample_index)
        id = self._encoded_ids[idx, CHUNK_ID_INDEX]
        chunk_index = idx

        if return_chunk_index:
            return id, chunk_index

        return id
