from __future__ import annotations

from uuid import UUID

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

from app.core.errors import AppError
from app.db.session import get_db
from app.models import Area, Dataset
from app.schemas.common import Envelope
from app.schemas.operations import VectorSelectionResponse
from app.schemas.selection_partitions import VectorPartitionSelectionRequest
from app.services.vector_feature_service import VectorFeatureService
from app.utils.response import envelope


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


def _product_identity(dataset: Dataset) -> str:
    metadata = dataset.source_metadata if isinstance(dataset.source_metadata, dict) else {}
    return str(metadata.get("product_key") or dataset.reference_layer_name or "")


@router.post(
    "/datasets/vector/partitions/select",
    response_model=Envelope[VectorSelectionResponse],
)
def select_vector_partitions(
    project_id: UUID,
    payload: VectorPartitionSelectionRequest,
    db: Session = Depends(get_db),
):
    datasets = db.query(Dataset).filter(Dataset.id.in_(payload.dataset_ids)).all()
    by_id = {dataset.id: dataset for dataset in datasets}
    ordered = [by_id.get(dataset_id) for dataset_id in payload.dataset_ids]
    if any(dataset is None or dataset.project_id != project_id for dataset in ordered):
        raise AppError(code="DATASET_NOT_FOUND", message="One or more selection partitions were not found", status_code=404)
    typed_datasets = [dataset for dataset in ordered if dataset is not None]
    if any(dataset.dataset_type not in {"vector", "geojson"} or dataset.status != "ready" for dataset in typed_datasets):
        raise AppError(
            code="INVALID_VECTOR_PARTITIONS",
            message="Every selection partition must be a ready vector dataset",
            status_code=409,
        )
    source_names = {dataset.source_name for dataset in typed_datasets}
    product_keys = {_product_identity(dataset) for dataset in typed_datasets}
    if len(source_names) != 1 or len(product_keys) != 1:
        raise AppError(
            code="VECTOR_PARTITION_SOURCE_MISMATCH",
            message="Selection partitions must belong to one governed source product",
            details={"source_names": sorted(str(value) for value in source_names), "product_keys": sorted(product_keys)},
            status_code=409,
        )

    selection_geometry = None
    selection_area_id = None
    if payload.area_id is not None:
        selection_area = db.get(Area, payload.area_id)
        if selection_area is None or selection_area.project_id != project_id:
            raise AppError(code="AREA_NOT_FOUND", message="Area not found", status_code=404)
        selection_geometry, _covers_full_area = VectorFeatureService.constrain_bbox_to_area(
            payload.bbox.model_dump(),
            selection_area.geometry,
        )
        selection_area_id = selection_area.id

    representative = typed_datasets[0]
    dataset_ids = [dataset.id for dataset in typed_datasets]
    result = VectorFeatureService.select_features_by_bbox(
        db,
        dataset_id=representative.id,
        dataset_ids=dataset_ids,
        bbox=payload.bbox.model_dump(),
        limit=payload.limit,
        dataset=representative,
        selection_geometry=selection_geometry,
        selection_area_id=selection_area_id,
        deduplicate_source_features=True,
    )
    result.update(
        partition_count=len(dataset_ids),
        source_name=representative.source_name,
        dataset_ids=dataset_ids,
    )
    return envelope(VectorSelectionResponse(**result).model_dump(exclude_none=True))