geointel/backend/app/services/aoi_operation_executor.py

from __future__ import annotations

from uuid import UUID

from geoalchemy2.shape import to_shape

from app.core.errors import AppError
from app.models import AoiOperation, AoiOperationPartition
from app.schemas.grb import GrbAcquireRequest
from app.schemas.dhmv import DhmvAcquireRequest
from app.schemas.spw_terrain import SpwTerrainAcquireRequest
from app.schemas.official_vector import OfficialVectorAcquireRequest
from app.schemas.flood_hazard import FloodHazardAcquireRequest
from app.schemas.thematic_raster import ThematicRasterAcquireRequest
from app.schemas.bathymetry import (
    BathymetryProfileAcquireRequest,
    MdkBathymetryAcquireRequest,
)
from app.schemas.job import JobCreate
from app.schemas.operations import VectorSelectionBBox
from app.schemas.orthophoto import OrthophotoAcquireRequest
from app.services.aoi_operation_service import AoiOperationService
from app.services.grb_acquisition_service import GrbAcquisitionService
from app.services.dhmv_acquisition_service import DhmvAcquisitionService
from app.services.spw_terrain_service import SpwTerrainService
from app.services.official_vector_acquisition_service import (
    OfficialVectorAcquisitionService,
)
from app.services.flood_hazard_acquisition_service import FloodHazardAcquisitionService
from app.services.thematic_raster_acquisition_service import (
    ThematicRasterAcquisitionService,
)
from app.services.walous_land_cover_service import WalousLandCoverService
from app.services.bathymetry_profile_acquisition_service import (
    BathymetryProfileAcquisitionService,
)
from app.services.mdk_bathymetry_acquisition_service import (
    MdkBathymetryAcquisitionService,
)
from app.services.job_service import JobService
from app.services.orthophoto_acquisition_service import OrthophotoAcquisitionService


class AoiOperationExecutor:
    """Execute one bounded partition through an existing governed provider."""

    @staticmethod
    def execute_next(db, project_id: UUID, operation_id: UUID) -> dict:
        partition = AoiOperationService.claim_next(db, project_id, operation_id)
        if partition is None:
            AoiOperationService._refresh_parent(db, operation_id)
            return AoiOperationService.read(db, project_id, operation_id)
        operation = db.get(AoiOperation, operation_id)
        child = JobService.create_job(
            db,
            JobCreate(
                job_type=f"aoi.{operation.operation_type}.partition",
                project_id=project_id,
                parameters_json={
                    "aoi_operation_id": str(operation_id),
                    "partition_id": str(partition.id),
                    "partition_key": partition.partition_key,
                    "provider_key": partition.provider_key,
                    "product_key": partition.product_key,
                },
            ),
        )
        partition = db.get(AoiOperationPartition, partition.id)
        partition.child_job_id = child.id
        db.add(partition)
        db.commit()
        JobService.mark_running(db, child.id)
        try:
            result = AoiOperationExecutor._dispatch(
                db, project_id, operation, partition
            )
            output_id = (
                result.get("output_dataset_id") if isinstance(result, dict) else None
            )
            JobService.mark_success(
                db,
                child.id,
                result=result,
                output_dataset_id=UUID(str(output_id)) if output_id else None,
            )
            return AoiOperationService.complete(
                db, project_id, operation_id, partition.id, result
            )
        except AppError as exc:
            JobService.mark_failed(
                db, child.id, exc.message, {"code": exc.code, "details": exc.details}
            )
            return AoiOperationService.fail(
                db,
                project_id,
                operation_id,
                partition.id,
                exc.message,
                AoiOperationExecutor._retryable(exc),
                {"code": exc.code, "details": exc.details},
            )
        except Exception:
            try:
                db.rollback()
                JobService.mark_failed(
                    db,
                    child.id,
                    "Unexpected partition execution error",
                    {"code": "AOI_PARTITION_INTERNAL_ERROR"},
                )
            finally:
                AoiOperationService.fail(
                    db,
                    project_id,
                    operation_id,
                    partition.id,
                    "Unexpected partition execution error",
                    True,
                    {"code": "AOI_PARTITION_INTERNAL_ERROR"},
                )
            raise

    @staticmethod
    def _dispatch(
        db, project_id: UUID, operation: AoiOperation, partition: AoiOperationPartition
    ) -> dict:
        geometry = to_shape(partition.geometry)
        min_x, min_y, max_x, max_y = geometry.bounds
        bbox = VectorSelectionBBox(
            min_x=min_x, min_y=min_y, max_x=max_x, max_y=max_y, crs="EPSG:4326"
        )
        force_refresh = bool(
            (operation.request_json or {})
            .get("parameters_json", {})
            .get("force_refresh", False)
        )
        if partition.provider_key == "grb":
            return GrbAcquisitionService.acquire(
                db,
                project_id,
                GrbAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "orthophoto":
            return OrthophotoAcquisitionService.acquire(
                db,
                project_id,
                OrthophotoAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "dhmv":
            return DhmvAcquisitionService.acquire(
                db,
                project_id,
                DhmvAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "spw_terrain":
            return SpwTerrainService.acquire(
                db,
                project_id,
                SpwTerrainAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "official_vector":
            return OfficialVectorAcquisitionService.acquire(
                db,
                project_id,
                OfficialVectorAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "flood_hazard":
            return FloodHazardAcquisitionService.acquire(
                db,
                project_id,
                FloodHazardAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "thematic_raster":
            return ThematicRasterAcquisitionService.acquire(
                db,
                project_id,
                ThematicRasterAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "walous":
            return WalousLandCoverService.acquire(
                db,
                project_id,
                ThematicRasterAcquireRequest(
                    bbox=bbox,
                    area_id=operation.area_id,
                    product_key=partition.product_key,
                    force_refresh=force_refresh,
                ),
            )
        if partition.provider_key == "bathymetry_profiles":
            return BathymetryProfileAcquisitionService.acquire(
                db,
                project_id,
                BathymetryProfileAcquireRequest(
                    bbox=bbox, area_id=operation.area_id, force_refresh=force_refresh
                ),
            )
        if partition.provider_key == "mdk_bathymetry":
            return MdkBathymetryAcquisitionService.acquire(
                db,
                project_id,
                MdkBathymetryAcquireRequest(
                    bbox=bbox, area_id=operation.area_id, force_refresh=force_refresh
                ),
            )
        raise AppError(
            code="AOI_PROVIDER_UNSUPPORTED",
            message="No governed AOI executor is registered for this provider",
            details={"provider_key": partition.provider_key},
            status_code=422,
        )

    @staticmethod
    def _retryable(error: AppError) -> bool:
        return error.status_code >= 500 or error.code.endswith(
            ("TIMEOUT", "UNAVAILABLE", "TLS_ERROR")
        )