from __future__ import annotations

from uuid import UUID

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

from app.db.session import get_db
from app.schemas.aoi_operation import AoiOperationCreate, AoiOperationList, AoiOperationRead, AoiPartitionCheckpoint, AoiPartitionComplete, AoiPartitionFail, AoiPartitionRead
from app.schemas.common import Envelope
from app.services.aoi_operation_service import AoiOperationService
from app.services.aoi_operation_executor import AoiOperationExecutor
from app.utils.response import envelope


router = APIRouter(prefix="/projects/{project_id}/aoi-operations", tags=["aoi-operations"])


@router.post("", status_code=201, response_model=Envelope[AoiOperationRead])
def create_operation(project_id: UUID, payload: AoiOperationCreate, db: Session = Depends(get_db)):
    return envelope(AoiOperationService.create(db, project_id, payload))


@router.get("", response_model=Envelope[AoiOperationList])
def list_operations(project_id: UUID, limit: int = Query(default=50, ge=1, le=200), db: Session = Depends(get_db)):
    return envelope(AoiOperationService.list(db, project_id, limit))


@router.get("/{operation_id}", response_model=Envelope[AoiOperationRead])
def read_operation(project_id: UUID, operation_id: UUID, db: Session = Depends(get_db)):
    return envelope(AoiOperationService.read(db, project_id, operation_id))


@router.post("/{operation_id}/partitions/claim", response_model=Envelope[AoiPartitionRead | None])
def claim_partition(project_id: UUID, operation_id: UUID, db: Session = Depends(get_db)):
    partition = AoiOperationService.claim_next(db, project_id, operation_id)
    return envelope(AoiPartitionRead.model_validate(partition).model_dump() if partition else None)


@router.post("/{operation_id}/execute-next", response_model=Envelope[AoiOperationRead])
def execute_next_partition(project_id: UUID, operation_id: UUID, db: Session = Depends(get_db)):
    return envelope(AoiOperationExecutor.execute_next(db, project_id, operation_id))


@router.put("/{operation_id}/partitions/{partition_id}/checkpoint", response_model=Envelope[AoiPartitionRead])
def checkpoint_partition(project_id: UUID, operation_id: UUID, partition_id: UUID, payload: AoiPartitionCheckpoint, db: Session = Depends(get_db)):
    partition = AoiOperationService.checkpoint(db, project_id, operation_id, partition_id, payload.checkpoint_json)
    return envelope(AoiPartitionRead.model_validate(partition).model_dump())


@router.post("/{operation_id}/partitions/{partition_id}/complete", response_model=Envelope[AoiOperationRead])
def complete_partition(project_id: UUID, operation_id: UUID, partition_id: UUID, payload: AoiPartitionComplete, db: Session = Depends(get_db)):
    return envelope(AoiOperationService.complete(db, project_id, operation_id, partition_id, payload.result_json, payload.skipped))


@router.post("/{operation_id}/partitions/{partition_id}/fail", response_model=Envelope[AoiOperationRead])
def fail_partition(project_id: UUID, operation_id: UUID, partition_id: UUID, payload: AoiPartitionFail, db: Session = Depends(get_db)):
    return envelope(AoiOperationService.fail(db, project_id, operation_id, partition_id, payload.error_message, payload.retryable, payload.details))