from __future__ import annotations

import warnings
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Self

import numpy as np
from async_tiff import TIFF
from async_tiff.enums import Compression as AsyncTiffCompression
from async_tiff.enums import (
    PhotometricInterpretation as AsyncTiffPhotometricInterpretation,
)
from async_tiff.enums import (
    PlanarConfiguration,
    SampleFormat,
)

from async_geotiff._colorinterp import infer_color_interpretation
from async_geotiff._crs import crs_from_geo_keys
from async_geotiff._fetch import FetchTileMixin
from async_geotiff._gdal_metadata import (
    BandStatistics,
    GDALMetadata,
    parse_gdal_metadata,
)
from async_geotiff._overview import Overview
from async_geotiff._read import ReadMixin
from async_geotiff._tile import TiledMixin
from async_geotiff._transform import TransformMixin, create_transform
from async_geotiff.colormap import Colormap
from async_geotiff.enums import (
    ColorInterp,
    Compression,
    Interleaving,
    PhotometricInterpretation,
)

if TYPE_CHECKING:
    from collections.abc import Mapping

    from affine import Affine
    from async_tiff import GeoKeyDirectory, ImageFileDirectory, ObspecInput
    from pyproj.crs import CRS


@dataclass(frozen=True, init=False, kw_only=True, repr=False)
class GeoTIFF(ReadMixin, FetchTileMixin, TiledMixin, TransformMixin):
    """A class representing a GeoTIFF image."""

    _crs: CRS | None = None
    """A cached CRS instance.

    We don't use functools.cached_property on the `crs` attribute because of typing
    issues.
    """

    _tiff: TIFF
    """The underlying async-tiff TIFF instance that we wrap.
    """

    _primary_ifd: ImageFileDirectory = field(init=False)
    """The primary (first) IFD of the GeoTIFF.

    Some tags, like most geo tags, only exist on the primary IFD.
    """

    _mask_ifd: ImageFileDirectory | None = None
    """The mask IFD of the full-resolution GeoTIFF, if any.
    """

    _gkd: GeoKeyDirectory = field(init=False)
    """The GeoKeyDirectory of the primary IFD.
    """

    _overviews: list[Overview] = field(init=False)
    """A list of overviews for the GeoTIFF.
    """

    _gdal_metadata: GDALMetadata | None = None
    """The metadata extracted from the GDALMetadata TIFF tag, if any."""

    @property
    def _alpha_band_idx(self) -> int | None:
        """The index of the alpha band, if any."""
        # TODO: when we support fetching partial bands, we need to check if the alpha
        # band is included in the bands we've fetched.
        # https://github.com/developmentseed/async-geotiff/issues/113
        alpha_band_idxs = [
            i
            for i, colorinterp in enumerate(self._geotiff.colorinterp)
            if colorinterp == ColorInterp.ALPHA
        ]
        if len(alpha_band_idxs) > 1:
            raise ValueError("Multiple alpha bands are not supported")

        return alpha_band_idxs[0] if alpha_band_idxs else None

    @property
    def _ifd(self) -> ImageFileDirectory:
        """An alias for the primary IFD to satisfy _fetch protocol."""
        return self._primary_ifd

    @property
    def _geotiff(self) -> GeoTIFF:
        """An alias for self to satisfy _fetch protocol."""
        return self

    def __init__(self, tiff: TIFF) -> None:
        """Create a GeoTIFF from an existing TIFF instance."""
        first_ifd = tiff.ifds[0]
        gkd = first_ifd.geo_key_directory

        # Validate that this is indeed a GeoTIFF
        if gkd is None:
            raise ValueError("TIFF does not contain GeoTIFF keys")

        if len(tiff.ifds) == 0:
            raise ValueError("TIFF does not contain any IFDs")

        # We use object.__setattr__ because the dataclass is frozen
        object.__setattr__(self, "_tiff", tiff)
        object.__setattr__(self, "_primary_ifd", first_ifd)
        object.__setattr__(self, "_gkd", gkd)
        object.__setattr__(
            self,
            "_gdal_metadata",
            parse_gdal_metadata(
                first_ifd.gdal_metadata,
                count=first_ifd.samples_per_pixel,
            ),
        )

        for ifd in tiff.ifds:
            if ifd.strip_byte_counts is not None or ifd.strip_offsets is not None:
                # This is not an error so that we can test metadata parsing for some
                # striped tiffs against rasterio.
                # Additionally some "single tile" tiffs are actually stored with strip
                # tags instead of tile tags
                warnings.warn(
                    "Striped GeoTIFFs are not supported.",
                    UserWarning,
                    stacklevel=2,
                )

        # Separate data IFDs and mask IFDs (skip the primary IFD at index 0)
        # Data IFDs are indexed by (width, height) for matching with masks
        data_ifds: dict[tuple[int, int], ImageFileDirectory] = {}
        mask_ifds: dict[tuple[int, int], ImageFileDirectory] = {}

        for ifd in tiff.ifds[1:]:
            dims = (ifd.image_width, ifd.image_height)
            if is_mask_ifd(ifd):
                mask_ifds[dims] = ifd
            else:
                data_ifds[dims] = ifd

        # Find and set the mask for the primary IFD (matches primary dimensions)
        if primary_mask_ifd := mask_ifds.get(
            (first_ifd.image_width, first_ifd.image_height),
        ):
            object.__setattr__(self, "_mask_ifd", primary_mask_ifd)

        # Build overviews, sorted by resolution (highest to lowest, i.e., largest first)
        # Sort by width * height descending
        sorted_dims = sorted(data_ifds.keys(), key=lambda d: d[0] * d[1], reverse=True)

        overviews: list[Overview] = []
        for dims in sorted_dims:
            data_ifd = data_ifds[dims]
            mask_ifd = mask_ifds.get(dims)

            ovr = Overview._create(  # noqa: SLF001
                geotiff=self,
                gkd=gkd,
                ifd=data_ifd,
                mask_ifd=mask_ifd,
            )
            overviews.append(ovr)

        object.__setattr__(self, "_overviews", overviews)

    @classmethod
    async def open(
        cls,
        path: str,
        *,
        store: ObspecInput,
        prefetch: int = 32768,
        multiplier: float = 2.0,
    ) -> Self:
        """Open a new GeoTIFF.

        Args:
            path: The path within the store to read from.
            store: The backend to use for data fetching.
            prefetch: The number of initial bytes to read up front.
            multiplier: The multiplier to use for readahead size growth. Must be
                greater than 1.0. For example, for a value of `2.0`, the first metadata
                read will be of size `prefetch`, and then the next read will be of size
                `prefetch * 2`.

        Returns:
            A GeoTIFF instance.

        """
        tiff = await TIFF.open(
            path=path,
            store=store,
            prefetch=prefetch,
            multiplier=multiplier,
        )
        return cls(tiff)

    @property
    def colorinterp(self) -> tuple[ColorInterp, ...]:
        """The color interpretation of each band in index order."""
        interps = list(
            infer_color_interpretation(
                count=self.count,
                photometric=self.photometric,
                extra_samples=self._primary_ifd.extra_samples or [],
            ),
        )

        if gdal_metadata := self._gdal_metadata:
            for band_index, color_interp in gdal_metadata.colorinterp.items():
                interps[band_index - 1] = color_interp

        return tuple(interps)

    @property
    def colormap(self) -> Colormap | None:
        """Return the Colormap stored in the file, if any.

        Returns:
            A Colormap instance if the dataset has a colormap, else None.

        """
        if upstream_colormap := self._primary_ifd.colormap:
            return Colormap(_cmap=upstream_colormap, _nodata=self.nodata)

        return None

    @property
    def compression(self) -> Compression | None:  # noqa: PLR0911
        """The compression algorithm used for the dataset.

        Returns None if the compression is not recognized.
        """
        match self._primary_ifd.compression:
            case AsyncTiffCompression.JPEG:
                return Compression.JPEG
            case AsyncTiffCompression.LZW:
                return Compression.LZW
            case AsyncTiffCompression.PackBits:
                return Compression.PACKBITS
            case AsyncTiffCompression.Deflate:
                return Compression.DEFLATE
            case AsyncTiffCompression.Uncompressed:
                return Compression.UNCOMPRESSED
            case AsyncTiffCompression.ZSTD:
                return Compression.ZSTD

        return None

    @property
    def count(self) -> int:
        """The number of raster bands in the full image."""
        return self._primary_ifd.samples_per_pixel

    @property
    def crs(self) -> CRS:
        """The dataset's coordinate reference system."""
        if self._crs is not None:
            return self._crs

        crs = crs_from_geo_keys(self._gkd)

        # We manually manage the cached property here because `@cached_property` messes
        # with type-checking
        object.__setattr__(self, "_crs", crs)

        return crs

    @property
    def dtype(self) -> np.dtype | None:  # noqa: PLR0911, C901
        """The NumPy data type of the image.

        Returns None if the data type is unknown/not supported.
        """
        formats = set(self._primary_ifd.sample_format)
        bits_per_sample = set(self._primary_ifd.bits_per_sample)

        if len(formats) != 1:
            raise ValueError("Mixed sample formats are not supported.")
        if len(bits_per_sample) != 1:
            raise ValueError("Mixed bits per sample are not supported.")

        match formats.pop(), bits_per_sample.pop():
            case (SampleFormat.Uint, 8):
                return np.dtype(np.uint8)
            case (SampleFormat.Uint, 16):
                return np.dtype(np.uint16)
            case (SampleFormat.Uint, 32):
                return np.dtype(np.uint32)
            case (SampleFormat.Uint, 64):
                return np.dtype(np.uint64)
            case (SampleFormat.Float, 32):
                return np.dtype(np.float32)
            case (SampleFormat.Float, 64):
                return np.dtype(np.float64)
            case (SampleFormat.Int, 8):
                return np.dtype(np.int8)
            case (SampleFormat.Int, 16):
                return np.dtype(np.int16)
            case (SampleFormat.Int, 32):
                return np.dtype(np.int32)
            case (SampleFormat.Int, 64):
                return np.dtype(np.int64)

        return None

    @property
    def height(self) -> int:
        """The height (number of rows) of the full image."""
        return self._primary_ifd.image_height

    @property
    def interleaving(self) -> Interleaving:
        """The interleaving scheme of the dataset."""
        planar_configuration = self._primary_ifd.planar_configuration

        if planar_configuration == PlanarConfiguration.Planar:
            return Interleaving.BAND

        return Interleaving.PIXEL

    @property
    def nodata(self) -> float | None:
        """The dataset's single nodata value."""
        nodata = self._primary_ifd.gdal_nodata
        if nodata is None:
            return None

        return float(nodata)

    @property
    def offsets(self) -> tuple[float, ...]:
        """The offset for each band, if any."""
        if gdal_metadata := self._gdal_metadata:
            return gdal_metadata.offsets

        return tuple([0.0] * self.count)

    @property
    def overviews(self) -> list[Overview]:
        """A list of overview levels for the dataset.

        Overviews are reduced-resolution versions of the main image used for faster
        rendering at lower zoom levels.

        This list of overviews is ordered from finest to coarsest resolution. The first
        element of the list is the highest-resolution after the base image.
        """
        return self._overviews

    @property
    def photometric(self) -> PhotometricInterpretation | None:  # noqa: PLR0911
        """The photometric interpretation of the dataset.

        Returns None if the photometric interpretation is not recognized.
        """
        match self._primary_ifd.photometric_interpretation:
            case AsyncTiffPhotometricInterpretation.WhiteIsZero:
                return PhotometricInterpretation.WHITE_IS_ZERO
            case AsyncTiffPhotometricInterpretation.BlackIsZero:
                return PhotometricInterpretation.BLACK_IS_ZERO
            case AsyncTiffPhotometricInterpretation.RGB:
                return PhotometricInterpretation.RGB
            case AsyncTiffPhotometricInterpretation.RGBPalette:
                return PhotometricInterpretation.RGBPALETTE
            case AsyncTiffPhotometricInterpretation.TransparencyMask:
                return PhotometricInterpretation.TRANSPARENCY_MASK
            case AsyncTiffPhotometricInterpretation.CMYK:
                return PhotometricInterpretation.CMYK
            case AsyncTiffPhotometricInterpretation.YCbCr:
                return PhotometricInterpretation.YCBCR
            case AsyncTiffPhotometricInterpretation.CIELab:
                return PhotometricInterpretation.CIELAB

        return None

    @property
    def scales(self) -> tuple[float, ...]:
        """The scale for each band, if any.

        The keys of the mapping are **1-based** band indices to match GDAL's convention.
        """
        if gdal_metadata := self._gdal_metadata:
            return gdal_metadata.scales

        return tuple([1.0] * self.count)

    @property
    def shape(self) -> tuple[int, int]:
        """Get the shape (height, width) of the full image."""
        return (self.height, self.width)

    @property
    def stored_stats(self) -> Mapping[int, BandStatistics] | None:
        """The **pre-existing** statistics for each GeoTIFF band, if available.

        These statistics are extracted from the GDALMetadata TIFF tag and are never
        computed on demand here.

        The keys of the mapping are **1-based** band indices to match GDAL's convention.
        """
        gdal_metadata = self._gdal_metadata

        # For an empty band_statistics dict, return None
        if gdal_metadata is not None and (stats := gdal_metadata.band_statistics):
            return MappingProxyType(stats)

        return None

    @property
    def tile_height(self) -> int:
        """The height in pixels per tile of the image."""
        # A TIFF with a single tile can omit tile_height
        return self._ifd.tile_height or self._ifd.image_height

    @property
    def tile_width(self) -> int:
        """The width in pixels per tile of the image."""
        # A TIFF with a single tile can omit tile_width
        return self._ifd.tile_width or self._ifd.image_width

    @property
    def transform(self) -> Affine:
        """Return the dataset's georeferencing transformation matrix.

        This transform maps pixel row/column coordinates to coordinates in the dataset's
        CRS.
        """
        return create_transform(
            model_tiepoint=self._primary_ifd.model_tiepoint,
            model_pixel_scale=self._primary_ifd.model_pixel_scale,
            model_transformation=self._primary_ifd.model_transformation,
            raster_type=self._gkd.raster_type,
        )

    @property
    def width(self) -> int:
        """The width (number of columns) of the full image."""
        return self._primary_ifd.image_width


def has_geokeys(ifd: ImageFileDirectory) -> bool:
    """Check if an IFD has GeoTIFF keys.

    Args:
        ifd: The IFD to check.

    """
    return ifd.geo_key_directory is not None


def is_mask_ifd(ifd: ImageFileDirectory) -> bool:
    """Check if an IFD is a mask IFD."""
    return (
        ifd.new_subfile_type is not None
        and ifd.new_subfile_type & 4 != 0
        and ifd.photometric_interpretation
        == AsyncTiffPhotometricInterpretation.TransparencyMask
    )
