"""Augmentation API routes."""

from fastapi import APIRouter, Query

from backend.web.core.auth import AdminTokenDep, DocumentRepoDep, DatasetRepoDep
from backend.web.schemas.admin.augmentation import (
    AugmentationBatchRequest,
    AugmentationBatchResponse,
    AugmentationConfigSchema,
    AugmentationPreviewRequest,
    AugmentationPreviewResponse,
    AugmentationTypeInfo,
    AugmentationTypesResponse,
    AugmentedDatasetListResponse,
    PresetInfo,
    PresetsResponse,
)


def register_augmentation_routes(router: APIRouter) -> None:
    """Register augmentation endpoints on the router."""

    @router.get(
        "/types",
        response_model=AugmentationTypesResponse,
        summary="List available augmentation types",
    )
    async def list_augmentation_types(
        admin_token: AdminTokenDep,
    ) -> AugmentationTypesResponse:
        """
        List all available augmentation types with descriptions and parameters.
        """
        from shared.augmentation.pipeline import (
            AUGMENTATION_REGISTRY,
            AugmentationPipeline,
        )

        types = []
        for name, aug_class in AUGMENTATION_REGISTRY.items():
            # Create instance with empty params to get preview params
            aug = aug_class({})
            types.append(
                AugmentationTypeInfo(
                    name=name,
                    description=(aug_class.__doc__ or "").strip(),
                    affects_geometry=aug_class.affects_geometry,
                    stage=AugmentationPipeline.STAGE_MAPPING[name],
                    default_params=aug.get_preview_params(),
                )
            )

        return AugmentationTypesResponse(augmentation_types=types)

    @router.get(
        "/presets",
        response_model=PresetsResponse,
        summary="Get augmentation presets",
    )
    async def get_presets(
        admin_token: AdminTokenDep,
    ) -> PresetsResponse:
        """Get predefined augmentation presets for common use cases."""
        from shared.augmentation.presets import list_presets

        presets = [PresetInfo(**p) for p in list_presets()]
        return PresetsResponse(presets=presets)

    @router.post(
        "/preview/{document_id}",
        response_model=AugmentationPreviewResponse,
        summary="Preview augmentation on document image",
    )
    async def preview_augmentation(
        document_id: str,
        request: AugmentationPreviewRequest,
        admin_token: AdminTokenDep,
        docs: DocumentRepoDep,
        page: int = Query(default=1, ge=1, description="Page number"),
    ) -> AugmentationPreviewResponse:
        """
        Preview a single augmentation on a document page.

        Returns URLs to original and augmented preview images.
        """
        from backend.web.services.augmentation_service import AugmentationService

        service = AugmentationService(doc_repo=docs)
        return await service.preview_single(
            document_id=document_id,
            page=page,
            augmentation_type=request.augmentation_type,
            params=request.params,
        )

    @router.post(
        "/preview-config/{document_id}",
        response_model=AugmentationPreviewResponse,
        summary="Preview full augmentation config on document",
    )
    async def preview_config(
        document_id: str,
        config: AugmentationConfigSchema,
        admin_token: AdminTokenDep,
        docs: DocumentRepoDep,
        page: int = Query(default=1, ge=1, description="Page number"),
    ) -> AugmentationPreviewResponse:
        """Preview complete augmentation pipeline on a document page."""
        from backend.web.services.augmentation_service import AugmentationService

        service = AugmentationService(doc_repo=docs)
        return await service.preview_config(
            document_id=document_id,
            page=page,
            config=config,
        )

    @router.post(
        "/batch",
        response_model=AugmentationBatchResponse,
        summary="Create augmented dataset (offline preprocessing)",
    )
    async def create_augmented_dataset(
        request: AugmentationBatchRequest,
        admin_token: AdminTokenDep,
        docs: DocumentRepoDep,
        datasets: DatasetRepoDep,
    ) -> AugmentationBatchResponse:
        """
        Create a new augmented dataset from an existing dataset.

        This runs as a background task. The augmented images are stored
        alongside the original dataset for training.
        """
        from backend.web.services.augmentation_service import AugmentationService

        service = AugmentationService(doc_repo=docs, dataset_repo=datasets)
        return await service.create_augmented_dataset(
            source_dataset_id=request.dataset_id,
            config=request.config,
            output_name=request.output_name,
            multiplier=request.multiplier,
        )

    @router.get(
        "/datasets",
        response_model=AugmentedDatasetListResponse,
        summary="List augmented datasets",
    )
    async def list_augmented_datasets(
        admin_token: AdminTokenDep,
        datasets: DatasetRepoDep,
        limit: int = Query(default=20, ge=1, le=100, description="Page size"),
        offset: int = Query(default=0, ge=0, description="Offset"),
    ) -> AugmentedDatasetListResponse:
        """List all augmented datasets."""
        from backend.web.services.augmentation_service import AugmentationService

        service = AugmentationService(dataset_repo=datasets)
        return await service.list_augmented_datasets(limit=limit, offset=offset)
