ModelForge-Public/backend/src/modelforge_api/api/routes/migrations.py

from __future__ import annotations

import uuid
from typing import Annotated

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

from modelforge_api.api.authorization import Admin, require_viewer
from modelforge_api.db import get_session
from modelforge_api.domain.migration_contracts import (
    BatchReport,
    CutoverOperationResponse,
    CutoverPrepare,
    CutoverReport,
    MigrationBatchResponse,
    MigrationEventResponse,
    MigrationPlanCreate,
    MigrationPlanResponse,
    MigrationValidationPolicyCreate,
    MigrationValidationPolicyResponse,
    PreflightReport,
    ReconciliationReport,
    RollbackReport,
    ShadowReport,
    StateAction,
    ValidationReport,
    ValidationSnapshotResponse,
)
from modelforge_api.services.migration_engine import MigrationEngineService

router = APIRouter(
    prefix="/api/v1/admin/migrations",
    tags=["migrations"],
    dependencies=[Depends(require_viewer)],
)


def get_service(session: Annotated[Session, Depends(get_session)]) -> MigrationEngineService:
    return MigrationEngineService(session)


Service = Annotated[MigrationEngineService, Depends(get_service)]


@router.get("/validation-policies", response_model=list[MigrationValidationPolicyResponse])
def validation_policies(service: Service, _admin: Admin) -> list[MigrationValidationPolicyResponse]:
    return service.validation_policies()


@router.post(
    "/validation-policies",
    response_model=MigrationValidationPolicyResponse,
    status_code=status.HTTP_201_CREATED,
)
def create_validation_policy(
    request: MigrationValidationPolicyCreate, service: Service, _admin: Admin
) -> MigrationValidationPolicyResponse:
    return service.create_validation_policy(request)


@router.get("/plans", response_model=list[MigrationPlanResponse])
def plans(
    service: Service,
    _admin: Admin,
    limit: Annotated[int, Query(ge=1, le=200)] = 200,
) -> list[MigrationPlanResponse]:
    return service.plans(limit=limit)


@router.post("/plans", response_model=MigrationPlanResponse, status_code=status.HTTP_201_CREATED)
def create_plan(
    request: MigrationPlanCreate, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.create_plan(request)


@router.get("/plans/{plan_id}", response_model=MigrationPlanResponse)
def plan(plan_id: uuid.UUID, service: Service, _admin: Admin) -> MigrationPlanResponse:
    return service.plan(plan_id)


@router.post("/plans/{plan_id}/preflight", response_model=MigrationPlanResponse)
def preflight(
    plan_id: uuid.UUID, request: PreflightReport, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.preflight(plan_id, request)


@router.post("/plans/{plan_id}/backfill/start", response_model=MigrationPlanResponse)
def start_backfill(
    plan_id: uuid.UUID, request: StateAction, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.start_backfill(plan_id, request)


@router.post("/plans/{plan_id}/backfill/pause", response_model=MigrationPlanResponse)
def pause_backfill(
    plan_id: uuid.UUID, request: StateAction, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.pause_backfill(plan_id, request)


@router.post("/plans/{plan_id}/batches", response_model=MigrationBatchResponse)
def record_batch(
    plan_id: uuid.UUID, request: BatchReport, service: Service, _admin: Admin
) -> MigrationBatchResponse:
    return service.record_batch(plan_id, request)


@router.get("/plans/{plan_id}/batches", response_model=list[MigrationBatchResponse])
def batches(plan_id: uuid.UUID, service: Service, _admin: Admin) -> list[MigrationBatchResponse]:
    return service.batches(plan_id)


@router.post("/plans/{plan_id}/validation", response_model=ValidationSnapshotResponse)
def validate_target(
    plan_id: uuid.UUID, request: ValidationReport, service: Service, _admin: Admin
) -> ValidationSnapshotResponse:
    return service.validate(plan_id, request)


@router.get("/plans/{plan_id}/validation", response_model=list[ValidationSnapshotResponse])
def validation_snapshots(
    plan_id: uuid.UUID, service: Service, _admin: Admin
) -> list[ValidationSnapshotResponse]:
    return service.validation_snapshots(plan_id)


@router.post("/plans/{plan_id}/shadow/start", response_model=MigrationPlanResponse)
def start_shadow(
    plan_id: uuid.UUID, request: StateAction, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.start_shadow(plan_id, request)


@router.post("/plans/{plan_id}/shadow/complete", response_model=MigrationPlanResponse)
def complete_shadow(
    plan_id: uuid.UUID, request: ShadowReport, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.complete_shadow(plan_id, request)


@router.post("/plans/{plan_id}/cutover/prepare", response_model=CutoverOperationResponse)
def prepare_cutover(
    plan_id: uuid.UUID, request: CutoverPrepare, service: Service, _admin: Admin
) -> CutoverOperationResponse:
    return service.prepare_cutover(plan_id, request)


@router.post("/plans/{plan_id}/cutover/report", response_model=CutoverOperationResponse)
def report_cutover(
    plan_id: uuid.UUID, request: CutoverReport, service: Service, _admin: Admin
) -> CutoverOperationResponse:
    return service.report_cutover(plan_id, request)


@router.post("/plans/{plan_id}/rollback", response_model=CutoverOperationResponse)
def rollback(
    plan_id: uuid.UUID, request: RollbackReport, service: Service, _admin: Admin
) -> CutoverOperationResponse:
    return service.rollback(plan_id, request)


@router.post("/reconcile", response_model=CutoverOperationResponse)
def reconcile(
    request: ReconciliationReport, service: Service, _admin: Admin
) -> CutoverOperationResponse:
    return service.reconcile(request)


@router.post("/plans/{plan_id}/cancel", response_model=MigrationPlanResponse)
def cancel(
    plan_id: uuid.UUID, request: StateAction, service: Service, _admin: Admin
) -> MigrationPlanResponse:
    return service.cancel(plan_id, request)


@router.get("/cutovers", response_model=list[CutoverOperationResponse])
def cutovers(
    service: Service,
    _admin: Admin,
    plan_id: Annotated[uuid.UUID | None, Query()] = None,
) -> list[CutoverOperationResponse]:
    return service.operations(plan_id)


@router.get("/events", response_model=list[MigrationEventResponse])
def events(
    service: Service,
    _admin: Admin,
    plan_id: Annotated[uuid.UUID | None, Query()] = None,
    limit: Annotated[int, Query(ge=1, le=1000)] = 200,
) -> list[MigrationEventResponse]:
    return service.events(plan_id, limit)