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.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


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 _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,
                    )
                clipped, clipped_transform = mask(source, [mapping(analysis_geometry)], crop=True, filled=False, indexes=[1])
                depth = np.ma.asarray(clipped[0], dtype="float64")
                raw = depth.filled(np.nan)
                selected_cells = geometry_mask([mapping(analysis_geometry)], out_shape=depth.shape, transform=clipped_transform, invert=True)
                nodata = source.nodata
                valid = selected_cells & ~np.ma.getmaskarray(depth) & np.isfinite(raw) & (raw > 0.0)
                if nodata is not None:
                    valid &= raw != float(nodata)
                values = raw[valid]
                selected_cell_count = int(selected_cells.sum())
                inundated_cell_count = int(values.size)
                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,
            )

        inundated_area_ha = inundated_cell_count * cell_area_m2 / 10_000.0
        metrics = [
            metric("modelled_inundated_area_ha", "Gemodelleerd overstroomd oppervlak", inundated_area_ha, "ha", "positive_depth_cells_times_cell_area"),
            metric(
                "modelled_inundated_share_pct",
                "Aandeel selectie met gemodelleerde diepte",
                inundated_cell_count / max(1, selected_cell_count) * 100.0,
                "%",
                "positive_depth_cells_divided_by_selected_cells",
            ),
        ]
        if 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=selected_cell_count,
            inundated_cell_count=inundated_cell_count,
            inundated_fraction=round(inundated_cell_count / max(1, selected_cell_count), 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,
            ),
            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,
        )
        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

        raw = partition.values
        valid = (
            partition.selected_cells
            & np.isfinite(raw)
            & (raw != FloodHazardAcquisitionService.NODATA)
            & (raw > 0.0)
        )
        values = raw[valid]
        selected_cell_count = int(partition.selected_cells.sum())
        inundated_cell_count = int(values.size)
        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,
            )

        inundated_area_ha = inundated_cell_count * cell_area_m2 / 10_000.0
        metrics = [
            metric(
                "modelled_inundated_area_ha",
                "Gemodelleerd overstroomd oppervlak",
                inundated_area_ha,
                "ha",
                "positive_depth_cells_times_cell_area",
            ),
            metric(
                "modelled_inundated_share_pct",
                "Aandeel selectie met gemodelleerde diepte",
                inundated_cell_count / max(1, selected_cell_count) * 100.0,
                "%",
                "positive_depth_cells_divided_by_selected_cells",
            ),
        ]
        if 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=selected_cell_count,
            inundated_cell_count=inundated_cell_count,
            inundated_fraction=round(inundated_cell_count / max(1, selected_cell_count), 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,
            ),
            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