import logging
import time
from typing import List
from typing import Optional

import h5py
import numpy
from numpy.typing import ArrayLike

from ..config import guess_dataset_config
from ..types import StrictPositiveIntegral
from ._base import DatasetWriterBase

logger = logging.getLogger(__name__)


class IndexedDatasetWriter(DatasetWriterBase):
    """Append arrays of the same shape to a new HDF5 dataset in a potentially non-sequential manner.

    Instead of creating a dataset with the :code:`h5py` API

    .. code-block::python

        h5group["mydataset"] = [[1,2,3], [4,5,6], [7,8,9]]

    it can be done like this

    .. code-block::python

        with IndexedDatasetWriter(h5group, "mydataset") as writer:
            writer.add_point(1, [4,5,6])
            writer.add_points([0, 2], [[1,2,3], [7,8,9]])

    Chunk size determination, compression and flushing is handled.
    """

    def __init__(
        self,
        parent: h5py.Group,
        name: str,
        npoints: Optional[StrictPositiveIntegral] = None,
        attrs: Optional[dict] = None,
        flush_period: Optional[float] = None,
        overwrite: bool = False,
    ) -> None:
        super().__init__(parent, name, attrs=attrs, flush_period=flush_period)
        self._npoints = npoints
        self._overwrite = overwrite
        self._npoints_added: int = 0

        self._buffer: List[tuple[int, ArrayLike]] = []

    def _create_dataset(self, first_data_point: numpy.ndarray) -> h5py.Dataset:
        scan_shape = (self._npoints,)
        detector_shape = first_data_point.shape
        dtype = first_data_point.dtype

        if self._npoints is None:
            shape = (1,) + detector_shape
            max_shape = scan_shape + detector_shape
        else:
            shape = scan_shape + detector_shape
            max_shape = None

        options = guess_dataset_config(
            scan_shape, detector_shape, dtype=dtype, max_shape=max_shape
        )
        options.update(
            shape=shape,
            dtype=dtype,
            fillvalue=numpy.nan,  # converts to 0 for integers
        )
        if max_shape:
            options["maxshape"] = max_shape

        if self._overwrite and self._name in self._parent:
            del self._parent[self._name]

        dset = self._parent.create_dataset(self._name, **options)
        if self._attrs:
            dset.attrs.update(self._attrs)
        return dset

    @property
    def npoints_added(self) -> int:
        return self._npoints_added

    def add_point(
        self,
        data: ArrayLike,
        index: int,
    ) -> bool:
        """Append one array to the dataset."""
        if self._dataset is None:
            self._dataset = self._create_dataset(numpy.asarray(data))

        self._buffer.append((index, data))
        self._npoints_added += 1
        return self.flush_buffer()

    def add_points(self, data: ArrayLike, indices: ArrayLike) -> bool:
        """Append several arrays at once to the dataset."""
        if self._dataset is None:
            self._dataset = self._create_dataset(numpy.asarray(data[0]))

        self._buffer.extend(zip(indices, data))
        self._npoints_added += len(data)
        return self.flush_buffer()

    def flush_buffer(self, align: bool = False) -> bool:
        if not self._buffer:
            return False

        if align:
            logger.debug(
                "Chunk-aligned flushing is only supported by DatasetWriter. "
                "Only flush when `flush_period` has passed since the last flush."
            )
            if self._flush_time_expired():
                return False

        indices, values = zip(*self._buffer)
        indices = numpy.asarray(indices)
        values = numpy.asarray(values)

        # Enlarge the dataset when needed
        max_index = int(indices.max())
        if max_index >= self._dataset.shape[0]:
            self._dataset.resize(max_index + 1, axis=0)

        # Copy data from buffer to HDF5
        self._dataset[indices] = values

        # Remove copied data from buffer
        self._buffer.clear()

        self._flush_hdf5()
        self._last_flush = time.time()
        return True
