import logging
import math
from contextlib import ExitStack
from pathlib import Path
from typing import Any

import rasterio
from pyproj import Transformer
from shapely.geometry import Point
from shapely.geometry.base import BaseGeometry

from flight_map.geo import Bounds

LOGGER = logging.getLogger(__name__)


class Elevation:
    """Metre-valued terrain for the separate map land-elevation summary."""

    def __init__(
        self,
        terrain: Path | None,
        reference: str = "unknown",
    ) -> None:
        self.terrain_path = terrain
        self.reference = reference
        self.stack = ExitStack()
        self.ground: Any = None
        self.ground_transform: Transformer | None = None
        self.land_samples = 0
        self.valid_land_samples = 0
        self.water_samples = 0

    @property
    def land_coverage(self) -> float:
        return self.valid_land_samples / self.land_samples if self.land_samples else 0.0

    def __enter__(self) -> "Elevation":
        with ExitStack() as stack:
            if self.terrain_path:
                self.ground = stack.enter_context(rasterio.open(self.terrain_path))
                self.ground_transform = self._transform(self.ground)
            self.stack = stack.pop_all()
        return self

    def __exit__(self, *args: object) -> None:
        self.stack.close()

    @staticmethod
    def _transform(dataset: Any) -> Transformer:
        if dataset.crs is None:
            raise ValueError("Elevation raster has no coordinate reference system.")
        if dataset.count < 1:
            raise ValueError("Elevation raster contains no data band.")
        return Transformer.from_crs("EPSG:4326", dataset.crs, always_xy=True)

    @staticmethod
    def sample(dataset: Any, transform: Transformer, lon: float, lat: float) -> float | None:
        x, y = transform.transform(lon, lat)
        row, col = dataset.index(x, y)
        if row < 0 or col < 0 or row >= dataset.height or col >= dataset.width:
            return None
        value = next(dataset.sample([(x, y)], indexes=1, masked=True))[0]
        if bool(getattr(value, "mask", False)):
            return None
        result = float(value) * dataset.scales[0] + dataset.offsets[0]
        return result if math.isfinite(result) else None

    def mean_ground(
        self, bounds: Bounds, samples: int = 64, *, water: BaseGeometry | None = None
    ) -> float | None:
        self.land_samples = self.valid_land_samples = self.water_samples = 0
        if not self.ground or self.ground_transform is None or self.reference == "unknown":
            return None
        if samples < 2:
            raise ValueError("Terrain sampling requires at least two samples per axis.")
        sin_south, sin_north = (
            math.sin(math.radians(bounds.south)),
            math.sin(math.radians(bounds.north)),
        )
        values = []
        # Uniform longitude and sin(latitude) samples represent equal spherical areas.
        for row in range(samples):
            lat = math.degrees(
                math.asin(sin_south + (sin_north - sin_south) * (row + 0.5) / samples)
            )
            for col in range(samples):
                lon = bounds.west + (bounds.east - bounds.west) * (col + 0.5) / samples
                if water is not None and water.covers(Point(lon, lat)):
                    self.water_samples += 1
                    continue
                self.land_samples += 1
                value = self.sample(self.ground, self.ground_transform, lon, lat)
                if value is None:
                    continue
                self.valid_land_samples += 1
                values.append(value)
        if not values:
            LOGGER.warning(
                "Average land elevation unavailable: no valid land samples "
                "(%d mapped water samples excluded).",
                self.water_samples,
            )
            return None
        if self.valid_land_samples < self.land_samples:
            LOGGER.warning(
                "Incomplete height map: %.1f%% of sampled land has elevation data "
                "(%d/%d); averaging valid land only. Missing land is not treated as water.",
                100 * self.land_coverage,
                self.valid_land_samples,
                self.land_samples,
            )
        LOGGER.info(
            "Land elevation coverage: %.1f%% (%d/%d); %d water samples excluded.",
            100 * self.land_coverage,
            self.valid_land_samples,
            self.land_samples,
            self.water_samples,
        )
        return sum(values) / len(values)