from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Self

import numpy as np
from async_tiff.enums import PlanarConfiguration
from numpy.ma import MaskedArray

from async_geotiff._transform import TransformMixin

if TYPE_CHECKING:
    from affine import Affine
    from async_tiff import Array
    from numpy.typing import NDArray
    from pyproj.crs import CRS

    from async_geotiff import GeoTIFF


@dataclass(frozen=True, kw_only=True, eq=False)
class RasterArray(TransformMixin):
    """Georeferenced array data from a GeoTIFF."""

    data: NDArray
    """The array data with shape (bands, height, width)."""

    mask: NDArray[np.bool_] | None
    """The mask array with shape (height, width), if any.

    Values of True indicate valid data; False indicates no data.
    """

    width: int
    """The width of the array in pixels."""

    height: int
    """The height of the array in pixels."""

    count: int
    """The number of bands in the array."""

    transform: Affine
    """The affine transform mapping pixel coordinates to geographic coordinates."""

    _alpha_band_idx: int | None
    """The index of the alpha band, if any.

    The alpha band lives in the `data` array but is checked in `as_masked`.
    """

    _geotiff: GeoTIFF
    """A reference to the parent GeoTIFF."""

    @classmethod
    def _create(  # noqa: PLR0913
        cls,
        *,
        data: Array,
        mask: Array | None,
        planar_configuration: PlanarConfiguration,
        transform: Affine,
        geotiff: GeoTIFF,
        alpha_band_idx: int | None,
    ) -> Self:
        """Create a RasterArray from async_tiff data.

        Handles axis reordering to ensure data is always in (bands, height, width)
        order, matching rasterio's convention.
        """
        data_arr = np.asarray(data, copy=False)
        if mask is not None:
            mask_arr = np.asarray(mask, copy=False).astype(np.bool_, copy=False)
            assert mask_arr.ndim == 3  # noqa: S101, PLR2004
            assert mask_arr.shape[2] == 1  # noqa: S101
            # This assumes it's always (height, width, 1)
            mask_arr = np.squeeze(mask_arr, axis=2)
        else:
            mask_arr = None

        assert data_arr.ndim == 3, f"Expected 3D array, got {data_arr.ndim}D"  # noqa: S101, PLR2004

        # async_tiff returns data in the native TIFF order:
        # - Chunky (pixel interleaved): (height, width, bands)
        # - Planar (band interleaved): (bands, height, width)
        # We always want (bands, height, width) to match rasterio.
        if planar_configuration == PlanarConfiguration.Chunky:
            # Transpose from (H, W, C) to (C, H, W)
            data_arr = np.moveaxis(data_arr, -1, 0)

        count, height, width = data_arr.shape

        return cls(
            data=data_arr,
            mask=mask_arr,
            width=width,
            height=height,
            count=count,
            transform=transform,
            _geotiff=geotiff,
            _alpha_band_idx=alpha_band_idx,
        )

    def as_masked(self) -> MaskedArray:
        """Return the data as a masked array using the array mask or nodata value.

        !!! warning
            In a numpy [`MaskedArray`][numpy.ma.MaskedArray], `True`
            indicates invalid (masked) data and `False` indicates valid data.

            This is the inverse convention of a GeoTIFF's mask. The boolean array
            [`RasterArray.mask`][async_geotiff.RasterArray.mask] uses `True` for valid
            data and `False` for invalid data.

        Returns:
            A masked array with the same shape as `data`, where invalid data
                (as indicated by the mask) is masked out.

        """
        if self.mask is not None:
            # MaskedArray constructor does not automatically broadcast inputs.
            # mask must have exactly the same shape as data.
            mask = np.broadcast_to(~self.mask, self.data.shape)
            return MaskedArray(self.data, mask=mask)

        if self.nodata is not None:
            return np.ma.masked_equal(self.data, self.nodata)

        if self._alpha_band_idx is not None:
            alpha_band = self.data[self._alpha_band_idx]
            single_band_mask = alpha_band == np.iinfo(alpha_band.dtype).min
            mask = np.broadcast_to(
                single_band_mask,
                (self.count - 1, self.height, self.width),
            )
            # Rasterio semantics say that the alpha band itself in `data` should be
            # considered valid data.
            mask = np.insert(mask, self._alpha_band_idx, False, axis=0)
            return MaskedArray(self.data, mask=mask)

        return MaskedArray(self.data)

    @property
    def crs(self) -> CRS:
        """The coordinate reference system of the array."""
        return self._geotiff.crs

    @property
    def nodata(self) -> float | None:
        """The nodata value for the array, if any."""
        return self._geotiff.nodata

    @property
    def shape(self) -> tuple[int, int, int]:
        """The shape of the data array as (bands, height, width)."""
        return self.data.shape
