from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path

import numpy as np
import rasterio
from rasterio.crs import CRS
from rasterio.warp import transform

from app.core.geo import PathPoint


class DemNotConfiguredError(NotImplementedError):
    """Raised when DEM access is requested before COG data is configured."""


def _dem_files(dem_path: Path) -> list[Path]:
    if not dem_path.exists():
        return []
    return sorted(
        path
        for pattern in ("*.tif", "*.tiff", "*.TIF", "*.TIFF")
        for path in dem_path.rglob(pattern)
        if path.is_file()
    )


def _point_in_bounds(x: float, y: float, bounds: object) -> bool:
    return bounds.left <= x <= bounds.right and bounds.bottom <= y <= bounds.top


def _to_dataset_crs(lat: float, lon: float, dst_crs: CRS | None) -> tuple[float, float]:
    if dst_crs is None or dst_crs == CRS.from_epsg(4326):
        return lon, lat
    xs, ys = transform(CRS.from_epsg(4326), dst_crs, [lon], [lat])
    return xs[0], ys[0]


def _sample_dataset(dataset: rasterio.io.DatasetReader, lat: float, lon: float) -> float | None:
    x, y = _to_dataset_crs(lat, lon, dataset.crs)
    if not _point_in_bounds(x, y, dataset.bounds):
        return None

    value = next(dataset.sample([(x, y)], masked=True))[0]
    if np.ma.is_masked(value):
        return None
    if dataset.nodata is not None and float(value) == float(dataset.nodata):
        return None
    if not np.isfinite(value):
        return None
    return float(value)


def elevation_at(
    lat: float,
    lon: float,
    surface: str = "dtm",
    dem_path: str | Path = "/data/dem",
) -> float:
    files = _dem_files(Path(dem_path))
    if not files:
        raise DemNotConfiguredError(f"DEM files are not found in {dem_path} for surface={surface}")

    for path in files:
        with rasterio.open(path) as dataset:
            value = _sample_dataset(dataset, lat, lon)
            if value is not None:
                return value

    raise DemNotConfiguredError(
        f"No DEM tile covers lat={lat}, lon={lon}, surface={surface}, dem_path={dem_path}"
    )


def elevations_along(
    points: Sequence[PathPoint],
    surface: str = "dtm",
    dem_path: str | Path = "/data/dem",
) -> np.ndarray:
    if not points:
        return np.array([], dtype=float)

    files = _dem_files(Path(dem_path))
    if not files:
        raise DemNotConfiguredError(f"DEM files are not found in {dem_path} for surface={surface}")

    values: list[float | None] = [None] * len(points)
    remaining = set(range(len(points)))

    for path in files:
        if not remaining:
            break
        with rasterio.open(path) as dataset:
            for index in list(remaining):
                point = points[index]
                value = _sample_dataset(dataset, point.lat, point.lon)
                if value is not None:
                    values[index] = value
                    remaining.remove(index)

    if remaining:
        missing = ", ".join(str(index) for index in sorted(remaining)[:10])
        raise DemNotConfiguredError(
            f"No DEM tile covers {len(remaining)} point(s), first missing indices: {missing}"
        )

    return np.array([float(value) for value in values], dtype=float)