from __future__ import annotations

import io
import math
from datetime import UTC, datetime
from pathlib import Path
from uuid import UUID

from geoalchemy2.shape import to_shape
from pyproj import Transformer
from shapely.geometry import box, mapping
from shapely.ops import transform as shapely_transform

from app.core.config import Settings, get_settings
from app.core.errors import AppError
from app.services.raster_cell_selection import select_cells
from app.models import Area, Dataset
from app.schemas.dhmv import (
    TerrainMetric,
    TerrainPartitionSelectionRequest,
    TerrainSelectionRequest,
    TerrainSelectionResponse,
    TerrainSelectionSummary,
)
from app.services.dhmv_acquisition_service import DhmvAcquisitionService
from app.services.raster_partition_analysis_service import (
    RasterPartitionAnalysisService,
)


class TerrainAnalysisService:
    SUPPORTED_PROVIDERS = {DhmvAcquisitionService.PROVIDER, "spw_terrain"}
    UNSUPPORTED_METRICS = ["water_depth_m", "water_volume_m3"]
    LIMITATION = (
        "Hoogte, reliëf en helling zijn afgeleid uit DHMV II. Afstroming vraagt bijkomende hydrologische modellering. "
        "Waterdiepte en watervolume zijn niet beschikbaar uit DTM/DSM alleen."
    )

    @staticmethod
    def _load_dataset(db, project_id: UUID, dataset_id: UUID) -> Dataset:
        dataset = db.get(Dataset, dataset_id)
        if not dataset or dataset.project_id != project_id:
            raise AppError(
                code="DATASET_NOT_FOUND", message="Dataset not found", status_code=404
            )
        if (
            dataset.dataset_type != "raster"
            or dataset.source_name not in TerrainAnalysisService.SUPPORTED_PROVIDERS
        ):
            raise AppError(
                code="INVALID_TERRAIN_DATASET",
                message="Terrain analysis requires a governed regional elevation raster",
                status_code=400,
            )
        if (
            dataset.status != "ready"
            or not dataset.storage_path
            or not Path(dataset.storage_path).is_file()
        ):
            raise AppError(
                code="DATASET_FILE_MISSING",
                message="Persisted terrain raster file is unavailable",
                status_code=404,
            )
        return dataset

    @staticmethod
    def _selection_geometry(db, project_id: UUID, payload: TerrainSelectionRequest):
        selection = box(
            payload.bbox.min_x,
            payload.bbox.min_y,
            payload.bbox.max_x,
            payload.bbox.max_y,
        )
        if payload.area_id is None:
            return selection
        area = db.get(Area, payload.area_id)
        if not area:
            raise AppError(
                code="AREA_NOT_FOUND", message="Area not found", status_code=404
            )
        if area.project_id != project_id:
            raise AppError(
                code="INVALID_DATASET_SCOPE",
                message="Area does not belong to this project",
                status_code=400,
            )
        selection = selection.intersection(to_shape(area.geometry))
        if selection.is_empty or selection.area <= 0:
            raise AppError(
                code="TERRAIN_SELECTION_OUTSIDE_AREA",
                message="Selection does not overlap the selected work area",
                status_code=422,
            )
        return selection

    @staticmethod
    def analyze(
        db,
        project_id: UUID,
        dataset_id: UUID,
        payload: TerrainSelectionRequest,
        *,
        settings: Settings | None = None,
    ) -> dict:
        resolved_settings = settings or get_settings()
        dataset = TerrainAnalysisService._load_dataset(db, project_id, dataset_id)
        selection_4326 = TerrainAnalysisService._selection_geometry(
            db, project_id, payload
        )
        try:
            import numpy as np
            import rasterio
            from rasterio.mask import mask
        except ImportError as exc:
            raise AppError(
                code="RASTER_PROCESSING_UNAVAILABLE",
                message="Rasterio and numpy are required for terrain analysis",
                status_code=503,
            ) from exc

        source_metadata = dataset.source_metadata or {}
        product_key = str(source_metadata.get("product_key") or "")
        surface_model = str(source_metadata.get("surface_model") or "")
        product_is_governed = (
            product_key in DhmvAcquisitionService._products()
            if dataset.source_name == DhmvAcquisitionService.PROVIDER
            else product_key == "spw_mnt_1m_2021_2022"
        )
        if not product_is_governed or surface_model not in {"terrain", "surface"}:
            raise AppError(
                code="INVALID_TERRAIN_METADATA",
                message="Regional terrain product provenance is incomplete",
                status_code=409,
            )
        vertical_unit_label = str(source_metadata.get("vertical_unit_label") or "m TAW")

        try:
            with rasterio.open(dataset.storage_path) as source:
                if source.crs is None:
                    raise AppError(
                        code="INVALID_DATASET_CRS",
                        message="Terrain raster CRS is missing",
                        status_code=409,
                    )
                transformer = Transformer.from_crs(
                    "EPSG:4326", source.crs, always_xy=True
                )
                selection_metric = shapely_transform(
                    transformer.transform, selection_4326
                )
                source_extent = box(*source.bounds)
                analysis_geometry = selection_metric.intersection(source_extent)
                if analysis_geometry.is_empty or analysis_geometry.area <= 0:
                    raise AppError(
                        code="TERRAIN_SELECTION_OUTSIDE_DATASET",
                        message="Selection does not overlap the persisted DHMV raster",
                        status_code=422,
                    )
                min_x, min_y, max_x, max_y = analysis_geometry.bounds
                expected_cells = math.ceil(
                    (max_x - min_x) / abs(source.res[0])
                ) * math.ceil((max_y - min_y) / abs(source.res[1]))
                if expected_cells > resolved_settings.dhmv_max_pixels:
                    raise AppError(
                        code="TERRAIN_SELECTION_TOO_LARGE",
                        message="Terrain analysis exceeds the configured raster cell limit",
                        details={
                            "pixel_count": expected_cells,
                            "max_pixels": resolved_settings.dhmv_max_pixels,
                        },
                        status_code=422,
                    )
                # ``all_touched`` keeps the values of cells the selection only
                # clips, so a selection finer than one cell still has data to
                # read. Which of those cells actually count is decided by
                # ``select_cells`` below, so the normal result is unchanged.
                clipped, clipped_transform = mask(
                    source,
                    [mapping(analysis_geometry)],
                    crop=True,
                    filled=False,
                    indexes=[1],
                    all_touched=True,
                )
                elevation = np.ma.asarray(clipped[0], dtype="float64")
                raw = elevation.filled(np.nan)
                nodata = source.nodata
                invalid = ~np.isfinite(raw)
                if nodata is not None:
                    invalid |= raw == float(nodata)
                cell_selection = select_cells(
                    analysis_geometry,
                    out_shape=elevation.shape,
                    transform=clipped_transform,
                    cell_area_m2=abs(float(source.res[0])) * abs(float(source.res[1])),
                )
                selected_cells = cell_selection.mask
                valid_mask = selected_cells & ~np.ma.getmaskarray(elevation) & ~invalid
                values = raw[valid_mask]
                if values.size == 0:
                    raise AppError(
                        code="TERRAIN_NO_VALID_DATA",
                        message="No valid terrain height cells occur in this selection",
                        status_code=422,
                    )

                resolution_x = abs(float(source.res[0]))
                resolution_y = abs(float(source.res[1]))
                slope_values = np.asarray([], dtype="float64")
                if raw.shape[0] >= 2 and raw.shape[1] >= 2:
                    surface = np.where(valid_mask, raw, np.nan)
                    gradient_y, gradient_x = np.gradient(
                        surface, resolution_y, resolution_x
                    )
                    slope = np.degrees(np.arctan(np.hypot(gradient_x, gradient_y)))
                    slope_values = slope[np.isfinite(slope) & valid_mask]
        except AppError:
            raise
        except Exception as exc:
            raise AppError(
                code="TERRAIN_ANALYSIS_FAILED",
                message="The persisted terrain raster could not be analysed",
                details={"reason": str(exc)},
                status_code=500,
            ) from exc

        def metric(
            key: str, label: str, value: float, unit: str, method: str
        ) -> TerrainMetric:
            return TerrainMetric(
                metric_key=key,
                metric_label=label,
                metric_value=round(float(value), 4),
                metric_unit=unit,
                aggregation_method=method,
            )

        prefix = "terrain" if surface_model == "terrain" else "surface"
        elevation_label = (
            "Gemiddelde maaiveldhoogte"
            if surface_model == "terrain"
            else "Gemiddelde oppervlaktehoogte"
        )
        metrics = [
            metric(
                f"{prefix}_elevation_mean_m",
                elevation_label,
                values.mean(),
                vertical_unit_label,
                "mean_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_min_m",
                "Laagste hoogte",
                values.min(),
                vertical_unit_label,
                "minimum_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_max_m",
                "Hoogste hoogte",
                values.max(),
                vertical_unit_label,
                "maximum_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_p10_m",
                "10e percentiel hoogte",
                np.percentile(values, 10),
                vertical_unit_label,
                "percentile_10_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_p90_m",
                "90e percentiel hoogte",
                np.percentile(values, 90),
                vertical_unit_label,
                "percentile_90_valid_cells",
            ),
            metric(
                "relief_m",
                "Reliëfverschil",
                values.max() - values.min(),
                "m",
                "maximum_minus_minimum",
            ),
        ]
        if slope_values.size:
            metrics.extend(
                [
                    metric(
                        "slope_mean_deg",
                        "Gemiddelde helling",
                        slope_values.mean(),
                        "°",
                        "mean_finite_gradient",
                    ),
                    metric(
                        "slope_p90_deg",
                        "90e percentiel helling",
                        np.percentile(slope_values, 90),
                        "°",
                        "percentile_90_finite_gradient",
                    ),
                    metric(
                        "slope_max_deg",
                        "Steilste helling",
                        slope_values.max(),
                        "°",
                        "maximum_finite_gradient",
                    ),
                ]
            )
        primary = metrics[0]
        selected_cell_count = int(selected_cells.sum())
        response = TerrainSelectionResponse(
            dataset_id=dataset.id,
            dataset_ids=[dataset.id],
            partition_count=1,
            product_key=product_key,
            surface_model=surface_model,
            selection_bbox=payload.bbox,
            selection_area_id=payload.area_id,
            sample_count=int(values.size),
            slope_sample_count=int(slope_values.size),
            coverage_ratio=round(float(values.size / max(1, selected_cell_count)), 6),
            cell_selection_warning=cell_selection.warning,
            resolution_m=round(max(resolution_x, resolution_y), 4),
            vertical_reference=str(
                source_metadata.get("vertical_reference")
                or DhmvAcquisitionService.VERTICAL_REFERENCE
            ),
            summary=TerrainSelectionSummary(
                metric_label=primary.metric_label,
                metric_value=primary.metric_value,
                metric_unit=primary.metric_unit,
                aggregation_method=primary.aggregation_method,
                primary_metric_key=primary.metric_key,
                metrics=metrics,
            ),
            unsupported_metrics=TerrainAnalysisService.UNSUPPORTED_METRICS,
            limitation_message=str(
                source_metadata.get("limitation_message")
                or TerrainAnalysisService.LIMITATION
            ),
            generated_at=datetime.now(UTC).isoformat(),
        )
        return response.model_dump(mode="json")

    @staticmethod
    def analyze_partitions(
        db,
        project_id: UUID,
        payload: TerrainPartitionSelectionRequest,
        *,
        settings: Settings | None = None,
    ) -> dict:
        resolved_settings = settings or get_settings()
        product = DhmvAcquisitionService._products().get(
            payload.product_key.strip().lower()
        )
        if product is None:
            raise AppError(
                code="DHMV_PRODUCT_NOT_SUPPORTED",
                message="Select a governed DHMV terrain or surface product",
                details={"product_key": payload.product_key},
                status_code=422,
            )
        selection_4326 = TerrainAnalysisService._selection_geometry(
            db, project_id, payload
        )
        partition = RasterPartitionAnalysisService.select(
            db,
            project_id,
            source_name=DhmvAcquisitionService.PROVIDER,
            product_key=product.key,
            selection_geometry_4326=selection_4326,
            nodata=DhmvAcquisitionService.NODATA,
            max_pixels=resolved_settings.dhmv_max_pixels,
            dataset_ids=payload.dataset_ids,
        )
        surface_models = {
            str((dataset.source_metadata or {}).get("surface_model") or "")
            for dataset in partition.datasets
        }
        if surface_models != {product.surface_model}:
            raise AppError(
                code="INVALID_TERRAIN_METADATA",
                message="DHMV partition provenance is incomplete",
                details={"surface_models": sorted(surface_models)},
                status_code=409,
            )

        try:
            import numpy as np
        except ImportError as exc:
            raise AppError(
                code="RASTER_PROCESSING_UNAVAILABLE",
                message="Numpy is required for partitioned terrain analysis",
                status_code=503,
            ) from exc

        raw = partition.values
        invalid = ~np.isfinite(raw) | (raw == DhmvAcquisitionService.NODATA)
        valid_mask = partition.selected_cells & ~invalid
        values = raw[valid_mask]
        if values.size == 0:
            raise AppError(
                code="TERRAIN_NO_VALID_DATA",
                message="No valid DHMV height cells occur in this selection",
                status_code=422,
            )
        slope_values = np.asarray([], dtype="float64")
        if raw.shape[0] >= 2 and raw.shape[1] >= 2:
            surface = np.where(valid_mask, raw, np.nan)
            gradient_y, gradient_x = np.gradient(
                surface,
                partition.resolution_y,
                partition.resolution_x,
            )
            slope = np.degrees(np.arctan(np.hypot(gradient_x, gradient_y)))
            slope_values = slope[np.isfinite(slope) & valid_mask]

        def metric(
            key: str, label: str, value: float, unit: str, method: str
        ) -> TerrainMetric:
            return TerrainMetric(
                metric_key=key,
                metric_label=label,
                metric_value=round(float(value), 4),
                metric_unit=unit,
                aggregation_method=method,
            )

        prefix = "terrain" if product.surface_model == "terrain" else "surface"
        elevation_label = (
            "Gemiddelde maaiveldhoogte"
            if product.surface_model == "terrain"
            else "Gemiddelde oppervlaktehoogte"
        )
        metrics = [
            metric(
                f"{prefix}_elevation_mean_m",
                elevation_label,
                values.mean(),
                "m TAW",
                "mean_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_min_m",
                "Laagste hoogte",
                values.min(),
                "m TAW",
                "minimum_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_max_m",
                "Hoogste hoogte",
                values.max(),
                "m TAW",
                "maximum_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_p10_m",
                "10e percentiel hoogte",
                np.percentile(values, 10),
                "m TAW",
                "percentile_10_valid_cells",
            ),
            metric(
                f"{prefix}_elevation_p90_m",
                "90e percentiel hoogte",
                np.percentile(values, 90),
                "m TAW",
                "percentile_90_valid_cells",
            ),
            metric(
                "relief_m",
                "Reliëfverschil",
                values.max() - values.min(),
                "m",
                "maximum_minus_minimum",
            ),
        ]
        if slope_values.size:
            metrics.extend(
                [
                    metric(
                        "slope_mean_deg",
                        "Gemiddelde helling",
                        slope_values.mean(),
                        "°",
                        "mean_finite_gradient",
                    ),
                    metric(
                        "slope_p90_deg",
                        "90e percentiel helling",
                        np.percentile(slope_values, 90),
                        "°",
                        "percentile_90_finite_gradient",
                    ),
                    metric(
                        "slope_max_deg",
                        "Steilste helling",
                        slope_values.max(),
                        "°",
                        "maximum_finite_gradient",
                    ),
                ]
            )
        primary = metrics[0]
        selected_cell_count = int(partition.selected_cells.sum())
        first_dataset = partition.datasets[0]
        response = TerrainSelectionResponse(
            dataset_id=first_dataset.id,
            dataset_ids=[dataset.id for dataset in partition.datasets],
            partition_count=len(partition.datasets),
            product_key=product.key,
            surface_model=product.surface_model,
            selection_bbox=payload.bbox,
            selection_area_id=payload.area_id,
            sample_count=int(values.size),
            slope_sample_count=int(slope_values.size),
            coverage_ratio=round(float(values.size / max(1, selected_cell_count)), 6),
            cell_selection_warning=partition.cell_selection_warning,
            resolution_m=round(max(partition.resolution_x, partition.resolution_y), 4),
            vertical_reference=DhmvAcquisitionService.VERTICAL_REFERENCE,
            summary=TerrainSelectionSummary(
                metric_label=primary.metric_label,
                metric_value=primary.metric_value,
                metric_unit=primary.metric_unit,
                aggregation_method=primary.aggregation_method,
                primary_metric_key=primary.metric_key,
                metrics=metrics,
            ),
            unsupported_metrics=TerrainAnalysisService.UNSUPPORTED_METRICS,
            limitation_message=(
                f"{TerrainAnalysisService.LIMITATION} De selectie werd exact berekend over "
                f"{len(partition.datasets)} persistente gemeentelijke rasterpartities."
            ),
            generated_at=datetime.now(UTC).isoformat(),
        )
        return response.model_dump(mode="json")

    @staticmethod
    def render_png(
        db, project_id: UUID, dataset_id: UUID, *, max_dimension: int = 1800
    ) -> bytes:
        dataset = TerrainAnalysisService._load_dataset(db, project_id, dataset_id)
        try:
            import numpy as np
            import rasterio
            from PIL import Image
            from rasterio.enums import Resampling
        except ImportError as exc:
            raise AppError(
                code="RASTER_PROCESSING_UNAVAILABLE",
                message="Rasterio, numpy and Pillow are required for terrain rendering",
                status_code=503,
            ) from exc

        try:
            with rasterio.open(dataset.storage_path) as source:
                scale = min(1.0, max_dimension / max(source.width, source.height))
                width = max(1, round(source.width * scale))
                height = max(1, round(source.height * scale))
                data = source.read(
                    1,
                    out_shape=(height, width),
                    masked=True,
                    resampling=Resampling.bilinear,
                )
                values = np.asarray(data.filled(np.nan), dtype="float64")
                valid = np.isfinite(values) & ~np.ma.getmaskarray(data)
                if not valid.any():
                    raise AppError(
                        code="TERRAIN_NO_VALID_DATA",
                        message="Terrain raster contains no renderable cells",
                        status_code=422,
                    )
                low, high = np.percentile(values[valid], [2, 98])
                if high <= low:
                    high = low + 1.0
                normalized = np.clip((values - low) / (high - low), 0.0, 1.0)
                stops = np.asarray([0.0, 0.25, 0.5, 0.75, 1.0])
                colors = np.asarray(
                    [
                        [30, 94, 91],
                        [79, 139, 102],
                        [194, 183, 105],
                        [173, 121, 79],
                        [105, 94, 108],
                    ],
                    dtype="float64",
                )
                rgba = np.zeros((height, width, 4), dtype="uint8")
                for channel in range(3):
                    rgba[:, :, channel] = np.interp(
                        normalized, stops, colors[:, channel]
                    ).astype("uint8")
                rgba[:, :, 3] = np.where(valid, 225, 0).astype("uint8")
                output = io.BytesIO()
                Image.fromarray(rgba).save(output, format="PNG", optimize=True)
                return output.getvalue()
        except AppError:
            raise
        except Exception as exc:
            raise AppError(
                code="TERRAIN_PREVIEW_FAILED",
                message="The persisted terrain raster could not be rendered",
                details={"reason": str(exc)},
                status_code=500,
            ) from exc