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 RasterNotConfiguredError(NotImplementedError):
    """Raised when raster files are missing or do not cover requested points."""


def raster_files(path: str | Path) -> list[Path]:
    raster_path = Path(path)
    if not raster_path.exists():
        return []
    return sorted(
        file_path
        for pattern in ("*.tif", "*.tiff", "*.TIF", "*.TIFF")
        for file_path in raster_path.rglob(pattern)
        if file_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 sample_at(lat: float, lon: float, path: str | Path, label: str) -> float:
    files = raster_files(path)
    if not files:
        raise RasterNotConfiguredError(f"{label} raster files are not found in {path}")

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

    raise RasterNotConfiguredError(f"No {label} raster tile covers lat={lat}, lon={lon}")


def sample_along(
    points: Sequence[PathPoint],
    path: str | Path,
    label: str,
    require_all: bool = True,
) -> np.ndarray:
    if not points:
        return np.array([], dtype=float)

    files = raster_files(path)
    if not files:
        raise RasterNotConfiguredError(f"{label} raster files are not found in {path}")

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

    for file_path in files:
        if not remaining:
            break
        with rasterio.open(file_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 and require_all:
        missing = ", ".join(str(index) for index in sorted(remaining)[:10])
        raise RasterNotConfiguredError(
            f"No {label} raster tile covers {len(remaining)} point(s), "
            f"first missing indices: {missing}"
        )

    return np.array([np.nan if value is None else float(value) for value in values], dtype=float)