geointel/backend/app/api/routes/detection.py

from __future__ import annotations

from uuid import UUID

from fastapi import APIRouter, Depends
from sqlalchemy.orm import Session

from app.db.session import get_db
from app.schemas import (
    AnalysisQaResponse,
    DetectionListResponse,
    DetectionModelsResponse,
    DetectionQaRequest,
    DetectionRead,
    DetectionRunListResponse,
    DetectionRunRead,
    DetectionRunRequest,
    DetectionRunResponse,
    Envelope,
    GeoJsonFeatureCollection,
    JobRead,
    ModelAssetListResponse,
    YoloPreflightResponse,
)
from app.services.detection_service import DetectionService
from app.services.model_asset_catalog_service import ModelAssetCatalogService
from app.services.model_registry_service import ModelRegistryService
from app.services.yolo_preflight_service import YoloPreflightService
from app.utils.response import envelope

router = APIRouter(prefix="/detection", tags=["detection"])


@router.get("/models", response_model=Envelope[DetectionModelsResponse])
def list_detection_models() -> dict:
    return envelope({"models": [model.model_dump() for model in ModelRegistryService.list_model_capabilities()]})


@router.get("/model-assets", response_model=Envelope[ModelAssetListResponse])
def list_detection_model_assets() -> dict:
    return envelope(ModelAssetCatalogService.list_assets().model_dump())


@router.get("/yolo/preflight", response_model=Envelope[YoloPreflightResponse])
def get_yolo_preflight(
    tile_manifest_path: str | None = None,
    check_model_load: bool = False,
    model_asset_id: str | None = None,
) -> dict:
    return envelope(
        YoloPreflightService.run(
            tile_manifest_path=tile_manifest_path,
            check_model_load=check_model_load,
            model_asset_id=model_asset_id,
        )
    )


@router.post("/run", response_model=Envelope[DetectionRunResponse])
def run_detection(payload: DetectionRunRequest, db: Session = Depends(get_db)) -> dict:
    result = DetectionService.run_detection(
        db=db,
        project_id=payload.project_id,
        dataset_id=payload.dataset_id,
        model_id=payload.model_id,
        model_asset_id=payload.model_asset_id,
        confidence_threshold=payload.confidence_threshold,
        class_filter=payload.class_filter,
        tile_manifest_path=payload.tile_manifest_path,
        parameters_json=payload.parameters_json,
    )
    return envelope(result.model_dump())


@router.post("/run-async", response_model=Envelope[JobRead])
def queue_detection(payload: DetectionRunRequest, db: Session = Depends(get_db)) -> dict:
    """Queue a detection run for the background worker.

    Tiled GPU inference takes minutes; ``POST /detection/run`` performs it
    inside the request and is only appropriate for a handful of tiles. Poll
    ``GET /jobs/{id}`` for the queued run instead.
    """

    job = DetectionService.enqueue_detection(
        db=db,
        project_id=payload.project_id,
        dataset_id=payload.dataset_id,
        model_id=payload.model_id,
        model_asset_id=payload.model_asset_id,
        confidence_threshold=payload.confidence_threshold,
        class_filter=payload.class_filter,
        tile_manifest_path=payload.tile_manifest_path,
        parameters_json=payload.parameters_json,
    )
    return envelope(JobRead.model_validate(job).model_dump(mode="json"))


@router.get("/runs", response_model=Envelope[DetectionRunListResponse])
def list_detection_runs(
    project_id: UUID | None = None,
    dataset_id: UUID | None = None,
    db: Session = Depends(get_db),
) -> dict:
    return envelope(DetectionService.list_runs(db, project_id=project_id, dataset_id=dataset_id).model_dump())


@router.get("/runs/{analysis_run_id}", response_model=Envelope[DetectionRunRead])
def get_detection_run(analysis_run_id: UUID, db: Session = Depends(get_db)) -> dict:
    return envelope(DetectionService.get_run(db, analysis_run_id).model_dump())


@router.get(
    "/runs/{analysis_run_id}/detections",
    response_model=Envelope[DetectionListResponse],
)
def list_detection_run_detections(
    analysis_run_id: UUID,
    dataset_id: UUID | None = None,
    class_name: str | None = None,
    min_confidence: float | None = None,
    db: Session = Depends(get_db),
) -> dict:
    return envelope(
        DetectionService.list_detections(
            db,
            analysis_run_id=analysis_run_id,
            dataset_id=dataset_id,
            class_name=class_name,
            min_confidence=min_confidence,
        ).model_dump()
    )


@router.get(
    "/datasets/{dataset_id}/detections",
    response_model=Envelope[DetectionListResponse],
)
def list_dataset_detections(
    dataset_id: UUID,
    analysis_run_id: UUID | None = None,
    class_name: str | None = None,
    min_confidence: float | None = None,
    db: Session = Depends(get_db),
) -> dict:
    return envelope(
        DetectionService.list_detections(
            db,
            analysis_run_id=analysis_run_id,
            dataset_id=dataset_id,
            class_name=class_name,
            min_confidence=min_confidence,
        ).model_dump()
    )


@router.get("/detections/{detection_id}", response_model=Envelope[DetectionRead])
def get_detection(detection_id: UUID, db: Session = Depends(get_db)) -> dict:
    return envelope(DetectionService.get_detection(db, detection_id).model_dump())


@router.get(
    "/runs/{analysis_run_id}/geojson",
    response_model=Envelope[GeoJsonFeatureCollection],
)
def get_detection_run_geojson(
    analysis_run_id: UUID,
    class_name: str | None = None,
    min_confidence: float | None = None,
    db: Session = Depends(get_db),
) -> dict:
    return envelope(
        DetectionService.detections_to_geojson(
            db,
            analysis_run_id=analysis_run_id,
            class_name=class_name,
            min_confidence=min_confidence,
        )
    )


@router.get(
    "/datasets/{dataset_id}/geojson",
    response_model=Envelope[GeoJsonFeatureCollection],
)
def get_dataset_detection_geojson(
    dataset_id: UUID,
    analysis_run_id: UUID | None = None,
    class_name: str | None = None,
    min_confidence: float | None = None,
    db: Session = Depends(get_db),
) -> dict:
    return envelope(
        DetectionService.detections_to_geojson(
            db,
            analysis_run_id=analysis_run_id,
            dataset_id=dataset_id,
            class_name=class_name,
            min_confidence=min_confidence,
        )
    )


@router.post(
    "/runs/{analysis_run_id}/qa/reference",
    response_model=Envelope[AnalysisQaResponse],
)
def compare_detection_run_with_reference(
    analysis_run_id: UUID,
    payload: DetectionQaRequest,
    db: Session = Depends(get_db),
) -> dict:
    return envelope(
        DetectionService.compare_detections_with_reference(
            db=db,
            analysis_run_id=analysis_run_id,
            reference_dataset_id=payload.reference_dataset_id,
            iou_threshold=payload.iou_threshold,
            class_name=payload.class_name,
            min_confidence=payload.min_confidence,
        )
    )