"""Training Dataset Endpoints."""

import logging
from typing import Annotated

from fastapi import APIRouter, HTTPException, Query

from backend.web.core.auth import (
    AdminTokenDep,
    DatasetRepoDep,
    DocumentRepoDep,
    AnnotationRepoDep,
    ModelVersionRepoDep,
    TrainingTaskRepoDep,
    FineTunePoolRepoDep,
)
from backend.web.schemas.admin import (
    DatasetCreateRequest,
    DatasetDetailResponse,
    DatasetDocumentItem,
    DatasetListItem,
    DatasetListResponse,
    DatasetResponse,
    DatasetTrainRequest,
    TrainingStatus,
    TrainingTaskResponse,
)
from backend.web.services.storage_helpers import get_storage_helper

from ._utils import _validate_uuid

logger = logging.getLogger(__name__)


def register_dataset_routes(router: APIRouter) -> None:
    """Register dataset endpoints on the router."""

    @router.post(
        "/datasets",
        response_model=DatasetResponse,
        summary="Create training dataset",
        description="Create a dataset from selected documents with train/val/test splits.",
    )
    async def create_dataset(
        request: DatasetCreateRequest,
        admin_token: AdminTokenDep,
        datasets: DatasetRepoDep,
        docs: DocumentRepoDep,
        annotations: AnnotationRepoDep,
    ) -> DatasetResponse:
        """Create a training dataset from document IDs."""
        from backend.web.services.dataset_builder import DatasetBuilder

        # Validate minimum document count for proper train/val/test split
        if len(request.document_ids) < 10:
            raise HTTPException(
                status_code=400,
                detail=f"Minimum 10 documents required for training dataset (got {len(request.document_ids)})",
            )

        dataset = datasets.create(
            name=request.name,
            description=request.description,
            train_ratio=request.train_ratio,
            val_ratio=request.val_ratio,
            seed=request.seed,
        )

        # Get storage paths from StorageHelper
        storage = get_storage_helper()
        datasets_dir = storage.get_datasets_base_path()
        admin_images_dir = storage.get_admin_images_base_path()

        if datasets_dir is None or admin_images_dir is None:
            raise HTTPException(
                status_code=500,
                detail="Storage not configured for local access",
            )

        builder = DatasetBuilder(
            datasets_repo=datasets,
            documents_repo=docs,
            annotations_repo=annotations,
            base_dir=datasets_dir,
        )
        try:
            builder.build_dataset(
                dataset_id=str(dataset.dataset_id),
                document_ids=request.document_ids,
                train_ratio=request.train_ratio,
                val_ratio=request.val_ratio,
                seed=request.seed,
                admin_images_dir=admin_images_dir,
            )
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))

        return DatasetResponse(
            dataset_id=str(dataset.dataset_id),
            name=dataset.name,
            status="ready",
            message="Dataset created successfully",
        )

    @router.get(
        "/datasets",
        response_model=DatasetListResponse,
        summary="List datasets",
    )
    async def list_datasets(
        admin_token: AdminTokenDep,
        datasets_repo: DatasetRepoDep,
        status: Annotated[str | None, Query(description="Filter by status")] = None,
        limit: Annotated[int, Query(ge=1, le=100)] = 20,
        offset: Annotated[int, Query(ge=0)] = 0,
    ) -> DatasetListResponse:
        """List training datasets."""
        datasets_list, total = datasets_repo.get_paginated(status=status, limit=limit, offset=offset)

        # Get active training tasks for each dataset (graceful degradation on error)
        dataset_ids = [str(d.dataset_id) for d in datasets_list]
        try:
            active_tasks = datasets_repo.get_active_training_tasks(dataset_ids)
        except Exception:
            logger.exception("Failed to fetch active training tasks")
            active_tasks = {}

        return DatasetListResponse(
            total=total,
            limit=limit,
            offset=offset,
            datasets=[
                DatasetListItem(
                    dataset_id=str(d.dataset_id),
                    name=d.name,
                    description=d.description,
                    status=d.status,
                    training_status=active_tasks.get(str(d.dataset_id), {}).get("status"),
                    active_training_task_id=active_tasks.get(str(d.dataset_id), {}).get("task_id"),
                    total_documents=d.total_documents,
                    total_images=d.total_images,
                    total_annotations=d.total_annotations,
                    created_at=d.created_at,
                )
                for d in datasets_list
            ],
        )

    @router.get(
        "/datasets/{dataset_id}",
        response_model=DatasetDetailResponse,
        summary="Get dataset detail",
    )
    async def get_dataset(
        dataset_id: str,
        admin_token: AdminTokenDep,
        datasets_repo: DatasetRepoDep,
    ) -> DatasetDetailResponse:
        """Get dataset details with document list."""
        _validate_uuid(dataset_id, "dataset_id")
        dataset = datasets_repo.get(dataset_id)
        if not dataset:
            raise HTTPException(status_code=404, detail="Dataset not found")

        docs = datasets_repo.get_documents(dataset_id)
        return DatasetDetailResponse(
            dataset_id=str(dataset.dataset_id),
            name=dataset.name,
            description=dataset.description,
            status=dataset.status,
            training_status=dataset.training_status,
            active_training_task_id=(
                str(dataset.active_training_task_id)
                if dataset.active_training_task_id
                else None
            ),
            train_ratio=dataset.train_ratio,
            val_ratio=dataset.val_ratio,
            seed=dataset.seed,
            total_documents=dataset.total_documents,
            total_images=dataset.total_images,
            total_annotations=dataset.total_annotations,
            dataset_path=dataset.dataset_path,
            error_message=dataset.error_message,
            documents=[
                DatasetDocumentItem(
                    document_id=str(d.document_id),
                    split=d.split,
                    page_count=d.page_count,
                    annotation_count=d.annotation_count,
                )
                for d in docs
            ],
            created_at=dataset.created_at,
            updated_at=dataset.updated_at,
        )

    @router.delete(
        "/datasets/{dataset_id}",
        summary="Delete dataset",
    )
    async def delete_dataset(
        dataset_id: str,
        admin_token: AdminTokenDep,
        datasets_repo: DatasetRepoDep,
    ) -> dict:
        """Delete a dataset and its files."""
        import shutil
        from pathlib import Path

        _validate_uuid(dataset_id, "dataset_id")
        dataset = datasets_repo.get(dataset_id)
        if not dataset:
            raise HTTPException(status_code=404, detail="Dataset not found")

        if dataset.dataset_path:
            dataset_dir = Path(dataset.dataset_path)
            if dataset_dir.exists():
                shutil.rmtree(dataset_dir)

        datasets_repo.delete(dataset_id)
        return {"message": "Dataset deleted"}

    @router.post(
        "/datasets/{dataset_id}/train",
        response_model=TrainingTaskResponse,
        summary="Start training from dataset",
        description="Create a training task. Set base_model_version_id in config for incremental training.",
    )
    async def train_from_dataset(
        dataset_id: str,
        request: DatasetTrainRequest,
        admin_token: AdminTokenDep,
        datasets_repo: DatasetRepoDep,
        models: ModelVersionRepoDep,
        tasks: TrainingTaskRepoDep,
        pool: FineTunePoolRepoDep,
    ) -> TrainingTaskResponse:
        """Create a training task from a dataset.

        For incremental training, set config.base_model_version_id to a model version UUID.
        The training will use that model as the starting point instead of a pretrained model.
        """
        _validate_uuid(dataset_id, "dataset_id")
        dataset = datasets_repo.get(dataset_id)
        if not dataset:
            raise HTTPException(status_code=404, detail="Dataset not found")
        if dataset.status != "ready":
            raise HTTPException(
                status_code=400,
                detail=f"Dataset is not ready (status: {dataset.status})",
            )

        config_dict = request.config.model_dump()

        # Resolve base_model_version_id to actual model path for incremental training
        base_model_version_id = config_dict.get("base_model_version_id")
        if base_model_version_id:
            _validate_uuid(base_model_version_id, "base_model_version_id")
            base_model = models.get(base_model_version_id)
            if not base_model:
                raise HTTPException(
                    status_code=404,
                    detail=f"Base model version not found: {base_model_version_id}",
                )
            # Chain prevention: never fine-tune from a fine-tuned model
            if getattr(base_model, "model_type", "base") == "finetune":
                original_base_id = getattr(base_model, "base_model_version_id", None)
                raise HTTPException(
                    status_code=400,
                    detail=(
                        f"Cannot chain fine-tunes. Model {base_model.version} is already "
                        f"a fine-tuned model. Select the original base model instead"
                        f"{f' (base_model_version_id: {original_base_id})' if original_base_id else ''}."
                    ),
                )
            # Pool threshold: require minimum verified pool entries for fine-tuning
            from backend.web.services.data_mixer import MIN_POOL_SIZE

            verified_count = pool.get_pool_count(verified_only=True)
            if verified_count < MIN_POOL_SIZE:
                raise HTTPException(
                    status_code=400,
                    detail=(
                        f"Fine-tuning requires at least {MIN_POOL_SIZE} verified pool entries "
                        f"(currently {verified_count}). Add more documents to the fine-tune "
                        f"pool and verify them before starting fine-tuning."
                    ),
                )

            # Store the resolved model path for the training worker
            config_dict["base_model_path"] = base_model.model_path
            config_dict["base_model_version"] = base_model.version
            logger.info(
                "Fine-tuning: using base model %s (%s) with %d verified pool entries",
                base_model.version,
                base_model.model_path,
                verified_count,
            )

        task_id = tasks.create(
            admin_token=admin_token,
            name=request.name,
            task_type="finetune" if base_model_version_id else "train",
            config=config_dict,
            dataset_id=str(dataset.dataset_id),
        )

        message = (
            f"Incremental training task created (base: v{config_dict.get('base_model_version', 'N/A')})"
            if base_model_version_id
            else "Training task created from dataset"
        )

        return TrainingTaskResponse(
            task_id=task_id,
            status=TrainingStatus.PENDING,
            message=message,
        )
