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

import rasterio
from pyproj import Transformer

from flight_map.geo import Bounds
from flight_map.storage import Observation


class Elevation:
    """Metre-valued terrain and optional h-H geoid separation rasters."""

    def __init__(
        self,
        terrain: Path | None,
        reference: str = "unknown",
        geoid: Path | None = None,
        confirm_geometric_ellipsoid: bool = False,
    ) -> None:
        self.terrain_path = terrain
        self.reference = reference
        self.geoid_path = geoid
        self.confirm_geometric_ellipsoid = confirm_geometric_ellipsoid
        self.stack = ExitStack()
        self.ground: Any = None
        self.geoid: Any = None
        self.ground_transform: Transformer | None = None
        self.geoid_transform: Transformer | None = None

    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)
            if self.geoid_path:
                self.geoid = stack.enter_context(rasterio.open(self.geoid_path))
                self.geoid_transform = self._transform(self.geoid)
            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) -> float | None:
        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
                value = self.sample(self.ground, self.ground_transform, lon, lat)
                if value is None:
                    return None
                values.append(value)
        return sum(values) / len(values)

    def aircraft_altitude(self, observation: Observation) -> float | None:
        if not self.confirm_geometric_ellipsoid or observation.geo_altitude is None:
            return None
        if self.reference == "wgs84-ellipsoid":
            return observation.geo_altitude
        if self.reference != "orthometric" or not self.geoid or self.geoid_transform is None:
            return None
        separation = self.sample(
            self.geoid, self.geoid_transform, observation.longitude, observation.latitude
        )
        if separation is None:
            return None
        return observation.geo_altitude - separation