from __future__ import annotations

import io
import math
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
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.flood_hazard import (
    FloodHazardMetric,
    FloodHazardPartitionSelectionRequest,
    FloodHazardSelectionRequest,
    FloodHazardSelectionResponse,
    FloodHazardSelectionSummary,
)
from app.services.flood_hazard_acquisition_service import FloodHazardAcquisitionService
from app.services.raster_partition_analysis_service import RasterPartitionAnalysisService


@dataclass(frozen=True)
class FloodHazardCellStatistics:
    """Cell populations behind one flood-hazard selection.

    Three populations, deliberately kept apart:

    ``selected``
        every cell whose centre falls inside the drawn selection;
    ``valid``
        the subset the VMM raster actually models — finite, not nodata;
    ``inundated``
        the subset of valid cells with a positive modelled depth.

    Risk is a share of what was modelled. Dividing by the selected cells
    instead silently reports "no data" as "no risk", which for a selection
    reaching past the modelled extent understates the hazard by whatever
    fraction of the rectangle the model never covered.
    """

    selected_cell_count: int
    valid_cell_count: int
    inundated_cell_count: int
    depth_values: Any

    @property
    def no_data_cell_count(self) -> int:
        return max(0, self.selected_cell_count - self.valid_cell_count)

    @property
    def data_coverage_ratio(self) -> float:
        if self.selected_cell_count <= 0:
            return 0.0
        return self.valid_cell_count / self.selected_cell_count

    @property
    def inundated_fraction(self) -> float | None:
        """``None`` when nothing was modelled: absence of data is not a zero."""

        if self.valid_cell_count <= 0:
            return None
        return self.inundated_cell_count / self.valid_cell_count

    def inundated_area_ha(self, cell_area_m2: float) -> float:
        return self.inundated_cell_count * cell_area_m2 / 10_000.0

    def analysed_area_ha(self, cell_area_m2: float) -> float:
        """Area the model actually covers inside the selection."""

        return self.valid_cell_count * cell_area_m2 / 10_000.0

    def selected_area_ha(self, cell_area_m2: float) -> float:
        """Area of the selection as rasterised, model coverage aside."""

        return self.selected_cell_count * cell_area_m2 / 10_000.0

    @classmethod
    def from_cells(cls, values: Any, selected: Any, *, nodata: float | None) -> "FloodHazardCellStatistics":
        import numpy as np

        raw = np.asarray(values, dtype="float64")
        selected_mask = np.asarray(selected, dtype=bool)

        has_data = selected_mask & np.isfinite(raw)
        if nodata is not None:
            has_data &= ~np.isclose(raw, float(nodata))
        # A modelled zero or negative depth is data: it says "dry here", which
        # is a different statement from "not modelled here".
        inundated = has_data & (raw > 0.0)

        return cls(
            selected_cell_count=int(selected_mask.sum()),
            valid_cell_count=int(has_data.sum()),
            inundated_cell_count=int(inundated.sum()),
            depth_values=raw[inundated],
        )


class FloodHazardAnalysisService:
    UNSUPPORTED_METRICS = [
        "bathymetry_depth_m",
        "permanent_water_volume_m3",
        "concurrent_flood_volume_m3",
    ]
    LIMITATION = (
        "Alle waarden horen bij het gekozen VMM-overstromingsscenario. De diepte-oppervlakte-integraal telt lokale "
        "gemodelleerde maxima op en is geen gelijktijdig opgeslagen watervolume, actuele waterstand of bathymetrie."
    )

    @staticmethod
    def _coverage_metrics(stats: "FloodHazardCellStatistics", cell_area_m2: float, metric) -> list[FloodHazardMetric]:
        """Headline metrics, each stating which population it is a share of.

        The analysed area is reported next to the drawn area so an operator can
        see immediately how much of the rectangle the flood model covers. A
        selection with no model data reports 0% coverage rather than 0% risk.
        """

        metrics = [
            metric(
                "modelled_inundated_area_ha",
                "Gemodelleerd overstroomd oppervlak",
                stats.inundated_area_ha(cell_area_m2),
                "ha",
                "positive_depth_cells_times_cell_area",
            ),
            metric(
                "modelled_inundated_share_pct",
                "Aandeel gemodelleerd gebied met diepte",
                0.0 if stats.inundated_fraction is None else stats.inundated_fraction * 100.0,
                "%",
                "positive_depth_cells_divided_by_modelled_cells",
            ),
            metric(
                "modelled_area_ha",
                "Oppervlak met overstromingsmodel",
                stats.analysed_area_ha(cell_area_m2),
                "ha",
                "modelled_cells_times_cell_area",
            ),
            metric(
                "selection_area_ha",
                "Oppervlak van de selectie",
                stats.selected_area_ha(cell_area_m2),
                "ha",
                "selected_cells_times_cell_area",
            ),
            metric(
                "model_coverage_pct",
                "Deel van de selectie met een model",
                stats.data_coverage_ratio * 100.0,
                "%",
                "modelled_cells_divided_by_selected_cells",
            ),
        ]
        return metrics

    @staticmethod
    def _combined_warning(stats: "FloodHazardCellStatistics", cell_selection_warning: str | None) -> str | None:
        parts = [
            part
            for part in (cell_selection_warning, FloodHazardAnalysisService._coverage_warning(stats))
            if part
        ]
        return " ".join(parts) if parts else None

    @staticmethod
    def _coverage_warning(stats: "FloodHazardCellStatistics") -> str | None:
        if stats.valid_cell_count <= 0:
            return (
                "Voor deze selectie bestaat geen VMM-overstromingsmodel. Er is dus geen overstromingsrisico "
                "gemeten; dit is geen bevestiging dat het risico nul is."
            )
        if stats.data_coverage_ratio < 0.999:
            return (
                f"Het VMM-model dekt {stats.data_coverage_ratio * 100:.1f}% van deze selectie. Percentages gelden "
                "voor het gemodelleerde deel, niet voor de volledige selectie."
            )
        return None

    @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 != FloodHazardAcquisitionService.PROVIDER:
            raise AppError(
                code="INVALID_FLOOD_HAZARD_DATASET",
                message="Flood-hazard analysis requires a governed VMM flood-depth 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 VMM flood-depth raster is unavailable", status_code=404)
        return dataset

    @staticmethod
    def _selection_geometry(db, project_id: UUID, payload: FloodHazardSelectionRequest):
        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)
        intersection = selection.intersection(to_shape(area.geometry))
        if intersection.is_empty or intersection.area <= 0:
            raise AppError(code="FLOOD_HAZARD_SELECTION_OUTSIDE_AREA", message="Selection does not overlap the selected work area", status_code=422)
        return intersection

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

        source_metadata = dataset.source_metadata or {}
        product_key = str(source_metadata.get("product_key") or "")
        product = FloodHazardAcquisitionService._products().get(product_key)
        if product is None or str(source_metadata.get("normalized_value_unit") or "") != "m":
            raise AppError(code="INVALID_FLOOD_HAZARD_METADATA", message="VMM flood-hazard provenance is incomplete", status_code=409)

        try:
            with rasterio.open(dataset.storage_path) as source:
                if source.crs is None:
                    raise AppError(code="INVALID_DATASET_CRS", message="VMM flood-depth 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)
                analysis_geometry = selection_metric.intersection(box(*source.bounds))
                if analysis_geometry.is_empty or analysis_geometry.area <= 0:
                    raise AppError(code="FLOOD_HAZARD_SELECTION_OUTSIDE_DATASET", message="Selection does not overlap the persisted flood-depth 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.flood_hazard_max_pixels:
                    raise AppError(
                        code="FLOOD_HAZARD_SELECTION_TOO_LARGE",
                        message="Flood-hazard analysis exceeds the configured raster cell limit",
                        details={"pixel_count": expected_cells, "max_pixels": resolved_settings.flood_hazard_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,
                )
                depth = np.ma.asarray(clipped[0], dtype="float64")
                raw = depth.filled(np.nan)
                cell_selection = select_cells(
                    analysis_geometry,
                    out_shape=depth.shape,
                    transform=clipped_transform,
                    cell_area_m2=abs(float(source.res[0])) * abs(float(source.res[1])),
                )
                selected_cells = cell_selection.mask
                # A masked cell carries no model value, so fold the mask into
                # the raw array before the populations are separated.
                raw = np.where(np.ma.getmaskarray(depth), np.nan, raw)
                stats = FloodHazardCellStatistics.from_cells(raw, selected_cells, nodata=source.nodata)
                values = stats.depth_values
                resolution_x = abs(float(source.res[0]))
                resolution_y = abs(float(source.res[1]))
                cell_area_m2 = resolution_x * resolution_y
        except AppError:
            raise
        except Exception as exc:
            raise AppError(
                code="FLOOD_HAZARD_ANALYSIS_FAILED",
                message="The persisted VMM flood-depth 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) -> FloodHazardMetric:
            return FloodHazardMetric(
                metric_key=key,
                metric_label=label,
                metric_value=round(float(value), 4),
                metric_unit=unit,
                aggregation_method=method,
            )

        metrics = FloodHazardAnalysisService._coverage_metrics(stats, cell_area_m2, metric)
        if stats.inundated_cell_count:
            metrics.extend(
                [
                    metric("modelled_depth_mean_m", "Gemiddelde gemodelleerde maximumdiepte", values.mean(), "m", "mean_positive_depth_cells"),
                    metric("modelled_depth_p90_m", "90e percentiel gemodelleerde maximumdiepte", np.percentile(values, 90), "m", "percentile_90_positive_depth_cells"),
                    metric("modelled_depth_max_m", "Hoogste gemodelleerde maximumdiepte", values.max(), "m", "maximum_positive_depth_cells"),
                    metric(
                        "modelled_max_depth_area_integral_m3",
                        "Diepte-oppervlakte-integraal (geen gelijktijdig volume)",
                        values.sum() * cell_area_m2,
                        "m3",
                        "sum_local_max_depth_times_cell_area",
                    ),
                ]
            )
        primary = metrics[0]
        response = FloodHazardSelectionResponse(
            dataset_id=dataset.id,
            dataset_ids=[dataset.id],
            partition_count=1,
            product_key=product.key,
            mechanism=product.mechanism,
            climate_context=product.climate_context,
            probability_class=product.probability_class,
            return_period_years=product.return_period_years,
            selection_bbox=payload.bbox,
            selection_area_id=payload.area_id,
            selected_cell_count=stats.selected_cell_count,
            valid_cell_count=stats.valid_cell_count,
            no_data_cell_count=stats.no_data_cell_count,
            data_coverage_ratio=round(stats.data_coverage_ratio, 6),
            inundated_cell_count=stats.inundated_cell_count,
            inundated_fraction=(
                None if stats.inundated_fraction is None else round(stats.inundated_fraction, 6)
            ),
            resolution_m=round(max(resolution_x, resolution_y), 4),
            summary=FloodHazardSelectionSummary(
                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,
            ),
            coverage_warning=FloodHazardAnalysisService._combined_warning(stats, cell_selection.warning),
            unsupported_metrics=FloodHazardAnalysisService.UNSUPPORTED_METRICS,
            limitation_message=FloodHazardAnalysisService.LIMITATION,
            generated_at=datetime.now(UTC).isoformat(),
        )
        return response.model_dump(mode="json")

    @staticmethod
    def analyze_partitions(
        db,
        project_id: UUID,
        payload: FloodHazardPartitionSelectionRequest,
        *,
        settings: Settings | None = None,
    ) -> dict:
        resolved_settings = settings or get_settings()
        product = FloodHazardAcquisitionService._products().get(payload.product_key.strip().lower())
        if product is None:
            raise AppError(
                code="FLOOD_HAZARD_PRODUCT_NOT_SUPPORTED",
                message="Select a governed VMM fluvial or pluvial flood-depth scenario",
                details={"product_key": payload.product_key},
                status_code=422,
            )
        selection_4326 = FloodHazardAnalysisService._selection_geometry(db, project_id, payload)
        partition = RasterPartitionAnalysisService.select(
            db,
            project_id,
            source_name=FloodHazardAcquisitionService.PROVIDER,
            product_key=product.key,
            selection_geometry_4326=selection_4326,
            nodata=FloodHazardAcquisitionService.NODATA,
            max_pixels=resolved_settings.flood_hazard_max_pixels,
            dataset_ids=payload.dataset_ids,
        )
        try:
            import numpy as np
        except ImportError as exc:
            raise AppError(
                code="RASTER_PROCESSING_UNAVAILABLE",
                message="Numpy is required for partitioned flood-hazard analysis",
                status_code=503,
            ) from exc

        stats = FloodHazardCellStatistics.from_cells(
            partition.values,
            partition.selected_cells,
            nodata=FloodHazardAcquisitionService.NODATA,
        )
        values = stats.depth_values
        cell_area_m2 = partition.resolution_x * partition.resolution_y

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

        metrics = FloodHazardAnalysisService._coverage_metrics(stats, cell_area_m2, metric)
        if stats.inundated_cell_count:
            metrics.extend(
                [
                    metric("modelled_depth_mean_m", "Gemiddelde gemodelleerde maximumdiepte", values.mean(), "m", "mean_positive_depth_cells"),
                    metric("modelled_depth_p90_m", "90e percentiel gemodelleerde maximumdiepte", np.percentile(values, 90), "m", "percentile_90_positive_depth_cells"),
                    metric("modelled_depth_max_m", "Hoogste gemodelleerde maximumdiepte", values.max(), "m", "maximum_positive_depth_cells"),
                    metric(
                        "modelled_max_depth_area_integral_m3",
                        "Diepte-oppervlakte-integraal (geen gelijktijdig volume)",
                        values.sum() * cell_area_m2,
                        "m3",
                        "sum_local_max_depth_times_cell_area",
                    ),
                ]
            )
        primary = metrics[0]
        first_dataset = partition.datasets[0]
        response = FloodHazardSelectionResponse(
            dataset_id=first_dataset.id,
            dataset_ids=[dataset.id for dataset in partition.datasets],
            partition_count=len(partition.datasets),
            product_key=product.key,
            mechanism=product.mechanism,
            climate_context=product.climate_context,
            probability_class=product.probability_class,
            return_period_years=product.return_period_years,
            selection_bbox=payload.bbox,
            selection_area_id=payload.area_id,
            selected_cell_count=stats.selected_cell_count,
            valid_cell_count=stats.valid_cell_count,
            no_data_cell_count=stats.no_data_cell_count,
            data_coverage_ratio=round(stats.data_coverage_ratio, 6),
            inundated_cell_count=stats.inundated_cell_count,
            inundated_fraction=(
                None if stats.inundated_fraction is None else round(stats.inundated_fraction, 6)
            ),
            resolution_m=round(max(partition.resolution_x, partition.resolution_y), 4),
            summary=FloodHazardSelectionSummary(
                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,
            ),
            coverage_warning=FloodHazardAnalysisService._combined_warning(stats, partition.cell_selection_warning),
            unsupported_metrics=FloodHazardAnalysisService.UNSUPPORTED_METRICS,
            limitation_message=(
                f"{FloodHazardAnalysisService.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 = FloodHazardAnalysisService._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 flood-hazard 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) & (values > 0.0)
                normalized = np.clip(values / 2.0, 0.0, 1.0)
                normalized = np.where(valid, normalized, 0.0)
                stops = np.asarray([0.0, 0.15, 0.35, 0.65, 1.0])
                colors = np.asarray(
                    [[190, 228, 255], [105, 184, 235], [42, 132, 201], [19, 83, 154], [8, 36, 92]],
                    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, np.clip(150 + normalized * 90, 0, 235), 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="FLOOD_HAZARD_PREVIEW_FAILED",
                message="The persisted VMM flood-depth raster could not be rendered",
                details={"reason": str(exc)},
                status_code=500,
            ) from exc