geointel/backend/app/services/aoi_operation_service.py

from __future__ import annotations

from collections import Counter
from datetime import datetime, timezone
from hashlib import sha256
import math
from uuid import UUID, uuid4

from geoalchemy2.shape import from_shape, to_shape
from pyproj import Transformer
from shapely.geometry import MultiPolygon, Polygon, box
from shapely.ops import transform

from app.core.errors import AppError
from app.core.config import get_settings
from app.models import AoiOperation, AoiOperationPartition, Area, Project
from app.schemas.aoi_operation import AoiOperationCreate


class AoiOperationService:
    MAX_PARTITIONS = 4096
    _to_metric = Transformer.from_crs("EPSG:4326", "EPSG:31370", always_xy=True)
    _to_wgs84 = Transformer.from_crs("EPSG:31370", "EPSG:4326", always_xy=True)
    SCOPE_AREA_NAMES = {
        "belgium": "Belgium land",
        "flanders": "Flanders",
        "wallonia": "Wallonia",
        "brussels": "Brussels-Capital Region",
        "belgian_north_sea": "Belgian part of the North Sea",
        "territorial_sea": "Belgian territorial sea (0-12 nautical miles)",
        "exclusive_economic_zone": "Belgian exclusive economic zone beyond territorial sea",
        "continental_shelf": "Belgian continental shelf beyond territorial sea",
    }

    @staticmethod
    def create(db, project_id: UUID, payload: AoiOperationCreate) -> dict:
        if db.get(Project, project_id) is None:
            raise AppError(
                code="PROJECT_NOT_FOUND", message="Project not found", status_code=404
            )
        geometry = AoiOperationService._resolve_geometry(db, project_id, payload)
        if payload.coverage_zone:
            geometry = AoiOperationService._clip_to_zone(
                db, project_id, geometry, payload.coverage_zone
            )
        geometry = AoiOperationService._as_multipolygon(geometry)
        metric_geometry = transform(AoiOperationService._to_metric.transform, geometry)
        partition_side_m = AoiOperationService._partition_side(
            payload.provider_key, payload.max_partition_side_m
        )
        cells = AoiOperationService._partition(metric_geometry, partition_side_m)
        operation_id = uuid4()
        now = datetime.now(timezone.utc)
        operation = AoiOperation(
            id=operation_id,
            project_id=project_id,
            area_id=payload.area_id,
            operation_type=payload.operation_type,
            status="queued",
            geometry=from_shape(geometry, srid=4326),
            request_json=payload.model_dump(mode="json", exclude_none=True),
            plan_json={
                "partition_strategy": "epsg31370_square_grid_intersection_v1",
                "max_partition_side_m": partition_side_m,
                "budget_source": "governed_provider_registry"
                if payload.max_partition_side_m is None
                else "stricter_operator_override",
                "partition_count": len(cells),
                "provider_key": payload.provider_key,
                "product_key": payload.product_key,
            },
            created_at=now,
        )
        db.add(operation)
        for ordinal, cell in enumerate(cells):
            wgs84 = transform(AoiOperationService._to_wgs84.transform, cell)
            wgs84 = AoiOperationService._as_multipolygon(wgs84)
            digest = sha256(wgs84.wkb).hexdigest()[:20]
            db.add(
                AoiOperationPartition(
                    id=uuid4(),
                    operation_id=operation_id,
                    partition_key=f"{payload.provider_key}:{payload.product_key}:{ordinal:05d}:{digest}",
                    provider_key=payload.provider_key,
                    product_key=payload.product_key,
                    ordinal=ordinal,
                    status="queued",
                    geometry=from_shape(wgs84, srid=4326),
                    attempt_count=0,
                    max_attempts=payload.max_attempts,
                    created_at=now,
                )
            )
        db.commit()
        return AoiOperationService.read(db, project_id, operation_id)

    @staticmethod
    def _clip_to_zone(db, project_id: UUID, geometry, zone: str):
        area_name = AoiOperationService.SCOPE_AREA_NAMES.get(zone)
        if area_name is None:
            raise AppError(
                code="AOI_COVERAGE_ZONE_UNSUPPORTED",
                message="Unknown governed coverage zone",
                details={"coverage_zone": zone},
                status_code=422,
            )
        scope = (
            db.query(Area)
            .filter(Area.project_id == project_id, Area.name == area_name)
            .first()
        )
        if scope is None:
            raise AppError(
                code="AOI_COVERAGE_ZONE_NOT_MATERIALIZED",
                message="The governed coverage-zone geometry is not persisted in this project",
                details={"coverage_zone": zone},
                status_code=409,
            )
        clipped = geometry.intersection(to_shape(scope.geometry))
        if clipped.is_empty:
            raise AppError(
                code="AOI_OUTSIDE_PROVIDER_ZONE",
                message="The AOI does not intersect the provider coverage zone",
                details={"coverage_zone": zone},
                status_code=422,
            )
        return clipped

    @staticmethod
    def _as_multipolygon(geometry) -> MultiPolygon:
        if isinstance(geometry, Polygon):
            return MultiPolygon([geometry])
        if isinstance(geometry, MultiPolygon):
            return geometry
        polygons = [
            part for part in getattr(geometry, "geoms", []) if isinstance(part, Polygon)
        ]
        if not polygons:
            raise AppError(
                code="AOI_GEOMETRY_EMPTY",
                message="AOI contains no polygonal area after clipping",
                status_code=422,
            )
        return MultiPolygon(polygons)

    @staticmethod
    def _partition_side(provider_key: str, requested: float | None) -> float:
        settings = get_settings()

        def raster_side(
            max_side_m: float, max_pixels: int, resolution_m: float
        ) -> float:
            # Keep every square grid cell within both the provider's spatial
            # extent limit and its decoded-pixel budget.  The small safety
            # margin absorbs ceil/edge rounding in the acquisition services.
            pixel_limited_side = (
                math.sqrt(float(max_pixels)) * float(resolution_m) * 0.99
            )
            return min(float(max_side_m), pixel_limited_side)

        budgets = {
            "orthophoto": float(settings.orthophoto_max_side_m),
            "grb": float(settings.grb_max_side_m),
            "dhmv": raster_side(
                settings.dhmv_max_side_m,
                settings.dhmv_max_pixels,
                settings.dhmv_resolution_m,
            ),
            "spw_terrain": raster_side(
                settings.spw_terrain_max_side_m,
                settings.spw_terrain_max_pixels,
                settings.spw_terrain_analysis_resolution_m,
            ),
            "official_vector": 20_000.0,
            "flood_hazard": raster_side(
                settings.flood_hazard_max_side_m,
                settings.flood_hazard_max_pixels,
                settings.flood_hazard_resolution_m,
            ),
            "thematic_raster": raster_side(
                settings.thematic_raster_max_side_m,
                settings.thematic_raster_max_pixels,
                10.0,
            ),
            "walous": raster_side(
                settings.walous_max_side_m,
                settings.walous_max_pixels,
                settings.walous_analysis_resolution_m,
            ),
            "bathymetry_profiles": 20_000.0,
            "mdk_bathymetry": 20_000.0,
        }
        if provider_key not in budgets:
            raise AppError(
                code="AOI_PROVIDER_UNSUPPORTED",
                message="No governed partition budget is registered for this provider",
                details={"provider_key": provider_key},
                status_code=422,
            )
        governed = budgets[provider_key]
        return min(governed, float(requested)) if requested is not None else governed

    @staticmethod
    def _resolve_geometry(db, project_id: UUID, payload: AoiOperationCreate):
        if (payload.area_id is None) == (payload.bbox is None):
            raise AppError(
                code="AOI_SELECTION_REQUIRED",
                message="Provide exactly one area_id or bbox",
                status_code=422,
            )
        if payload.area_id is not None:
            area = db.get(Area, payload.area_id)
            if area is None or area.project_id != project_id:
                raise AppError(
                    code="AREA_NOT_FOUND", message="Area not found", status_code=404
                )
            return to_shape(area.geometry)
        bbox = payload.bbox
        if bbox is None or bbox.crs != "EPSG:4326":
            raise AppError(
                code="INVALID_AOI_CRS",
                message="AOI bbox must use EPSG:4326",
                status_code=422,
            )
        return box(bbox.min_x, bbox.min_y, bbox.max_x, bbox.max_y)

    @staticmethod
    def _partition(geometry, side_m: float) -> list:
        min_x, min_y, max_x, max_y = geometry.bounds
        columns = max(1, math.ceil((max_x - min_x) / side_m))
        rows = max(1, math.ceil((max_y - min_y) / side_m))
        if columns * rows > AoiOperationService.MAX_PARTITIONS:
            raise AppError(
                code="AOI_PARTITION_LIMIT_EXCEEDED",
                message="AOI requires too many bounded partitions",
                details={
                    "candidate_count": columns * rows,
                    "max_partitions": AoiOperationService.MAX_PARTITIONS,
                },
                status_code=422,
            )
        partitions = []
        for row in range(rows):
            for column in range(columns):
                clipped = geometry.intersection(
                    box(
                        min_x + column * side_m,
                        min_y + row * side_m,
                        min(min_x + (column + 1) * side_m, max_x),
                        min(min_y + (row + 1) * side_m, max_y),
                    )
                )
                if not clipped.is_empty and clipped.area > 0:
                    partitions.append(clipped)
        return partitions

    @staticmethod
    def read(db, project_id: UUID, operation_id: UUID) -> dict:
        operation = db.get(AoiOperation, operation_id)
        if operation is None or operation.project_id != project_id:
            raise AppError(
                code="AOI_OPERATION_NOT_FOUND",
                message="AOI operation not found",
                status_code=404,
            )
        partitions = (
            db.query(AoiOperationPartition)
            .filter(AoiOperationPartition.operation_id == operation_id)
            .order_by(AoiOperationPartition.ordinal)
            .all()
        )
        counts = Counter(partition.status for partition in partitions)
        complete = counts["success"] + counts["skipped"]
        return {
            "id": operation.id,
            "project_id": operation.project_id,
            "area_id": operation.area_id,
            "parent_job_id": operation.parent_job_id,
            "operation_type": operation.operation_type,
            "status": operation.status,
            "request_json": operation.request_json,
            "plan_json": operation.plan_json,
            "result_json": operation.result_json,
            "error_message": operation.error_message,
            "progress": round(complete / len(partitions), 6) if partitions else 0.0,
            "partition_counts": dict(counts),
            "partitions": partitions,
            "created_at": operation.created_at,
            "started_at": operation.started_at,
            "finished_at": operation.finished_at,
        }

    @staticmethod
    def list(db, project_id: UUID, limit: int = 50) -> dict:
        rows = (
            db.query(AoiOperation)
            .filter(AoiOperation.project_id == project_id)
            .order_by(AoiOperation.created_at.desc())
            .limit(limit)
            .all()
        )
        return {
            "items": [AoiOperationService.read(db, project_id, row.id) for row in rows],
            "total": len(rows),
        }

    @staticmethod
    def claim_next(db, project_id: UUID, operation_id: UUID):
        operation = db.get(AoiOperation, operation_id)
        if operation is None or operation.project_id != project_id:
            raise AppError(
                code="AOI_OPERATION_NOT_FOUND",
                message="AOI operation not found",
                status_code=404,
            )
        partition = (
            db.query(AoiOperationPartition)
            .filter(
                AoiOperationPartition.operation_id == operation_id,
                AoiOperationPartition.status == "queued",
            )
            .order_by(AoiOperationPartition.ordinal)
            .with_for_update(skip_locked=True)
            .first()
        )
        if partition is None:
            return None
        now = datetime.now(timezone.utc)
        partition.status = "running"
        partition.started_at = now
        partition.attempt_count += 1
        partition.error_message = None
        operation.status = "running"
        operation.started_at = operation.started_at or now
        db.add(partition)
        db.add(operation)
        db.commit()
        db.refresh(partition)
        return partition

    @staticmethod
    def checkpoint(
        db, project_id: UUID, operation_id: UUID, partition_id: UUID, checkpoint: dict
    ):
        partition = AoiOperationService._partition_row(
            db, project_id, operation_id, partition_id
        )
        if partition.status != "running":
            raise AppError(
                code="AOI_PARTITION_NOT_RUNNING",
                message="Only a running partition can be checkpointed",
                status_code=409,
            )
        partition.checkpoint_json = checkpoint
        db.add(partition)
        db.commit()
        db.refresh(partition)
        return partition

    @staticmethod
    def complete(
        db,
        project_id: UUID,
        operation_id: UUID,
        partition_id: UUID,
        result: dict,
        skipped: bool = False,
    ):
        partition = AoiOperationService._partition_row(
            db, project_id, operation_id, partition_id
        )
        if partition.status == "success" or partition.status == "skipped":
            return AoiOperationService.read(db, project_id, operation_id)
        if partition.status != "running":
            raise AppError(
                code="AOI_PARTITION_NOT_RUNNING",
                message="Only a running partition can complete",
                status_code=409,
            )
        partition.status = "skipped" if skipped else "success"
        partition.result_json = result
        partition.finished_at = datetime.now(timezone.utc)
        db.add(partition)
        db.commit()
        AoiOperationService._refresh_parent(db, operation_id)
        return AoiOperationService.read(db, project_id, operation_id)

    @staticmethod
    def fail(
        db,
        project_id: UUID,
        operation_id: UUID,
        partition_id: UUID,
        message: str,
        retryable: bool,
        details: dict,
    ):
        partition = AoiOperationService._partition_row(
            db, project_id, operation_id, partition_id
        )
        partition.error_message = message
        partition.result_json = {"details": details}
        partition.status = (
            "queued"
            if retryable and partition.attempt_count < partition.max_attempts
            else "failed"
        )
        partition.finished_at = (
            None if partition.status == "queued" else datetime.now(timezone.utc)
        )
        db.add(partition)
        db.commit()
        AoiOperationService._refresh_parent(db, operation_id)
        return AoiOperationService.read(db, project_id, operation_id)

    @staticmethod
    def _partition_row(db, project_id, operation_id, partition_id):
        operation = db.get(AoiOperation, operation_id)
        partition = db.get(AoiOperationPartition, partition_id)
        if (
            operation is None
            or operation.project_id != project_id
            or partition is None
            or partition.operation_id != operation_id
        ):
            raise AppError(
                code="AOI_PARTITION_NOT_FOUND",
                message="AOI partition not found",
                status_code=404,
            )
        return partition

    @staticmethod
    def _refresh_parent(db, operation_id):
        operation = db.get(AoiOperation, operation_id)
        partitions = (
            db.query(AoiOperationPartition)
            .filter(AoiOperationPartition.operation_id == operation_id)
            .order_by(AoiOperationPartition.ordinal)
            .all()
        )
        statuses = [partition.status for partition in partitions]
        output_dataset_ids = []
        for partition in partitions:
            output_id = (
                (partition.result_json or {}).get("output_dataset_id")
                if isinstance(partition.result_json, dict)
                else None
            )
            if output_id and str(output_id) not in output_dataset_ids:
                output_dataset_ids.append(str(output_id))
        operation.result_json = {
            "partition_count": len(partitions),
            "completed_partition_count": sum(
                status in {"success", "skipped"} for status in statuses
            ),
            "failed_partition_count": statuses.count("failed"),
            "output_dataset_ids": output_dataset_ids,
            "merge_contract": "source_aware_spatial_union",
            "vector_deduplication": "source_feature_id_then_geometry",
            "raster_deduplication": "governed_mosaic_grid",
            "complete_coverage": bool(statuses)
            and all(status in {"success", "skipped"} for status in statuses),
        }
        now = datetime.now(timezone.utc)
        if statuses and all(status in {"success", "skipped"} for status in statuses):
            operation.status = "success"
            operation.finished_at = now
            operation.error_message = None
        elif "failed" in statuses and not any(
            status in {"queued", "running"} for status in statuses
        ):
            operation.status = (
                "partial"
                if any(status in {"success", "skipped"} for status in statuses)
                else "failed"
            )
            operation.finished_at = now
            operation.error_message = "One or more bounded source partitions failed; inspect partition evidence."
        db.add(operation)
        db.commit()