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()