from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass, field

import defusedxml.ElementTree as ET  # noqa: N817

from async_geotiff.enums import ColorInterp


@dataclass
class BandStatistics:
    """Statistics for a single band in a GeoTIFF."""

    max: float | None = None
    """The maximum pixel value in the band."""
    min: float | None = None
    """The minimum pixel value in the band."""
    mean: float | None = None
    """The mean pixel value in the band."""
    std: float | None = None
    """The standard deviation of the pixel values in the band."""
    valid_percent: float | None = None
    """The percentage of valid pixels in the band."""


@dataclass(frozen=True)
class GDALMetadata:
    """Metadata extracted from the GDALMetadata TIFF tag."""

    band_statistics: dict[int, BandStatistics]
    """A mapping of band index to statistics for that band.

    This band index is **1-based**.
    """

    offsets: tuple[float, ...]

    scales: tuple[float, ...]

    colorinterp: dict[int, ColorInterp] = field(default_factory=dict)
    """A mapping of 1-based band index to overridden ColorInterp.

    When present, these values override what would otherwise be inferred
    from the photometric interpretation tag.
    """


LOWER_CASE_COLORINTERP_MAPPING: dict[str, ColorInterp] = {}
for color in ColorInterp:
    LOWER_CASE_COLORINTERP_MAPPING[color.name.lower()] = color


def parse_gdal_metadata(  # noqa: C901
    gdal_metadata: str | None,
    *,
    count: int,
) -> GDALMetadata | None:
    if gdal_metadata is None:
        return None

    root = ET.fromstring(gdal_metadata)

    if root.tag != "GDALMetadata":
        raise ValueError("Not a GDALMetadata XML block")

    band_statistics: defaultdict[int, BandStatistics] = defaultdict(BandStatistics)
    offsets: list[float] = [0.0] * count
    scales: list[float] = [1.0] * count
    colorinterp: dict[int, ColorInterp] = {}

    for elem in root.findall("Item"):
        name = elem.attrib.get("name")
        sample = elem.attrib.get("sample")
        role = elem.attrib.get("role")
        text = elem.text or ""
        match name:
            # Add 1 to get a 1-based band index to match GDAL.
            case "STATISTICS_MAXIMUM" if sample is not None:
                band_statistics[int(sample) + 1].max = float(text)
            case "STATISTICS_MEAN" if sample is not None:
                band_statistics[int(sample) + 1].mean = float(text)
            case "STATISTICS_MINIMUM" if sample is not None:
                band_statistics[int(sample) + 1].min = float(text)
            case "STATISTICS_STDDEV" if sample is not None:
                band_statistics[int(sample) + 1].std = float(text)
            case "STATISTICS_VALID_PERCENT" if sample is not None:
                band_statistics[int(sample) + 1].valid_percent = float(text)
            case "OFFSET" if sample is not None:
                offsets[int(sample)] = float(text)
            case "SCALE" if sample is not None:
                scales[int(sample)] = float(text)
            case "COLORINTERP" if role == "colorinterp" and sample is not None:
                colorinterp[int(sample) + 1] = LOWER_CASE_COLORINTERP_MAPPING[
                    text.lower()
                ]

    return GDALMetadata(
        band_statistics=dict(band_statistics),
        offsets=tuple(offsets),
        scales=tuple(scales),
        colorinterp=colorinterp,
    )
