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

from __future__ import annotations

from uuid import UUID

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

from app.api.guest_scope import (
    assert_guest_project_scope,
    guest_project_scope,
    guest_scoped_project_filter,
)
from app.db.session import get_db
from app.schemas import (
    AnalysisQaResponse,
    Envelope,
    GeoJsonFeatureCollection,
    JobRead,
    SegmentationListResponse,
    SegmentationModelsResponse,
    SegmentationQaRequest,
    SegmentationRead,
    SegmentationRunListResponse,
    SegmentationRunRead,
    SegmentationRunRequest,
    SegmentationRunResponse,
)
from app.services.model_registry_service import ModelRegistryService
from app.services.dataset_service import DatasetService
from app.services.detection_service import DetectionService
from app.services.segmentation_service import SegmentationService
from app.utils.response import envelope

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


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


@router.post("/run", response_model=Envelope[SegmentationRunResponse])
def run_segmentation(
    payload: SegmentationRunRequest,
    request: Request,
    db: Session = Depends(get_db),
) -> dict:
    assert_guest_project_scope(request, payload.project_id)
    result = SegmentationService.run_segmentation(
        db=db,
        project_id=payload.project_id,
        dataset_id=payload.dataset_id,
        model_id=payload.model_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_segmentation(
    payload: SegmentationRunRequest,
    request: Request,
    db: Session = Depends(get_db),
) -> dict:
    """Queue a segmentation run for the background worker.

    Configured segmentation walks the same tile manifest as detection and is
    just as unsuited to running inside the request. Poll ``GET /jobs/{id}``.
    """

    assert_guest_project_scope(request, payload.project_id)
    job = SegmentationService.enqueue_segmentation(
        db=db,
        project_id=payload.project_id,
        dataset_id=payload.dataset_id,
        model_id=payload.model_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[SegmentationRunListResponse])
def list_segmentation_runs(
    request: Request,
    project_id: UUID | None = None,
    dataset_id: UUID | None = None,
    limit: int = Query(default=DetectionService.DEFAULT_RUN_LIST_LIMIT, ge=0, le=5_000),
    offset: int = Query(default=0, ge=0),
    db: Session = Depends(get_db),
) -> dict:
    project_id = guest_scoped_project_filter(request, project_id)
    return envelope(
        SegmentationService.list_runs(
            db, project_id=project_id, dataset_id=dataset_id, limit=limit, offset=offset
        ).model_dump()
    )


@router.get("/runs/{analysis_run_id}", response_model=Envelope[SegmentationRunRead])
def get_segmentation_run(
    analysis_run_id: UUID,
    request: Request,
    db: Session = Depends(get_db),
) -> dict:
    run = SegmentationService.get_run(db, analysis_run_id)
    assert_guest_project_scope(request, run.project_id)
    return envelope(run.model_dump())


@router.get(
    "/runs/{analysis_run_id}/segmentations",
    response_model=Envelope[SegmentationListResponse],
)
def list_segmentation_run_outputs(
    analysis_run_id: UUID,
    request: Request,
    dataset_id: UUID | None = None,
    class_name: str | None = None,
    min_confidence: float | None = None,
    limit: int = Query(
        default=DetectionService.DEFAULT_RESULT_LIMIT,
        ge=0,
        le=50_000,
        description="Maximum results to return; 0 returns everything. Highest confidence first.",
    ),
    offset: int = Query(default=0, ge=0),
    db: Session = Depends(get_db),
) -> dict:
    if guest_project_scope(request) is not None:
        run = SegmentationService.get_run(db, analysis_run_id)
        assert_guest_project_scope(request, run.project_id)
    return envelope(
        SegmentationService.list_segmentations(
            db,
            limit=limit,
            offset=offset,
            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}/segmentations",
    response_model=Envelope[SegmentationListResponse],
)
def list_dataset_segmentations(
    dataset_id: UUID,
    request: Request,
    analysis_run_id: UUID | None = None,
    class_name: str | None = None,
    min_confidence: float | None = None,
    limit: int = Query(
        default=DetectionService.DEFAULT_RESULT_LIMIT,
        ge=0,
        le=50_000,
        description="Maximum results to return; 0 returns everything. Highest confidence first.",
    ),
    offset: int = Query(default=0, ge=0),
    db: Session = Depends(get_db),
) -> dict:
    if guest_project_scope(request) is not None:
        dataset = DatasetService.get_dataset(db, dataset_id)
        assert_guest_project_scope(request, dataset.project_id)
    return envelope(
        SegmentationService.list_segmentations(
            db,
            limit=limit,
            offset=offset,
            analysis_run_id=analysis_run_id,
            dataset_id=dataset_id,
            class_name=class_name,
            min_confidence=min_confidence,
        ).model_dump()
    )


@router.get("/segmentations/{segmentation_id}", response_model=Envelope[SegmentationRead])
def get_segmentation(
    segmentation_id: UUID,
    request: Request,
    db: Session = Depends(get_db),
) -> dict:
    segmentation = SegmentationService.get_segmentation(db, segmentation_id)
    assert_guest_project_scope(request, segmentation.project_id)
    return envelope(segmentation.model_dump())


@router.get(
    "/runs/{analysis_run_id}/geojson",
    response_model=Envelope[GeoJsonFeatureCollection],
)
def get_segmentation_run_geojson(
    analysis_run_id: UUID,
    request: Request,
    class_name: str | None = None,
    min_confidence: float | None = None,
    limit: int = Query(
        default=DetectionService.DEFAULT_RESULT_LIMIT,
        ge=0,
        le=50_000,
        description="Maximum results to return; 0 returns everything. Highest confidence first.",
    ),
    db: Session = Depends(get_db),
) -> dict:
    if guest_project_scope(request) is not None:
        run = SegmentationService.get_run(db, analysis_run_id)
        assert_guest_project_scope(request, run.project_id)
    return envelope(
        SegmentationService.segmentations_to_geojson(
            db,
            limit=limit,
            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_segmentation_geojson(
    dataset_id: UUID,
    request: Request,
    analysis_run_id: UUID | None = None,
    class_name: str | None = None,
    min_confidence: float | None = None,
    limit: int = Query(
        default=DetectionService.DEFAULT_RESULT_LIMIT,
        ge=0,
        le=50_000,
        description="Maximum results to return; 0 returns everything. Highest confidence first.",
    ),
    db: Session = Depends(get_db),
) -> dict:
    if guest_project_scope(request) is not None:
        dataset = DatasetService.get_dataset(db, dataset_id)
        assert_guest_project_scope(request, dataset.project_id)
    return envelope(
        SegmentationService.segmentations_to_geojson(
            db,
            limit=limit,
            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_segmentation_run_with_reference(
    analysis_run_id: UUID,
    payload: SegmentationQaRequest,
    request: Request,
    db: Session = Depends(get_db),
) -> dict:
    if guest_project_scope(request) is not None:
        run = SegmentationService.get_run(db, analysis_run_id)
        assert_guest_project_scope(request, run.project_id)
    return envelope(
        SegmentationService.compare_segmentations_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,
            calibration_thresholds=payload.calibration_thresholds,
        )
    )