backend/app/api/datasets.py

"""Admin dataset management endpoints (auth-protected, under /admin).

Covers list/view/purge, image→dataset lookup, embed-all, and match-all. The dataset *load* job
endpoints live in api/jobs.py (background job + polling for the dashboard's load form).
"""
from __future__ import annotations

from datetime import date

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import select
from sqlalchemy.orm import Session

from ..db import get_db
from ..models import Dataset, Picture, User
from ..models.base import SubjectType
from ..schemas.dataset import (
    DatasetDetailOut,
    DatasetOut,
    EmbedAllResult,
    ImageDatasetOut,
    MatchDatasetResult,
    PurgeResult,
)
from ..security import require_admin
from ..services import datasets as ds

router = APIRouter(prefix="/admin", tags=["admin-datasets"])


def _get_dataset(db: Session, dataset_id: int) -> Dataset:
    dataset = db.get(Dataset, dataset_id)
    if not dataset:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "Dataset not found")
    return dataset


@router.get("/datasets", response_model=list[DatasetOut])
def list_datasets(
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> list[DatasetOut]:
    datasets = db.execute(select(Dataset).order_by(Dataset.creation_time.desc())).scalars().all()
    out: list[DatasetOut] = []
    for d in datasets:
        item = DatasetOut.model_validate(d)
        item.case_count = ds.case_count(db, d)
        out.append(item)
    return out


@router.get("/datasets/{dataset_id}", response_model=DatasetDetailOut)
def get_dataset(
    dataset_id: int,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> DatasetDetailOut:
    dataset = _get_dataset(db, dataset_id)
    stats = ds.dataset_stats(db, dataset)
    base = DatasetOut.model_validate(dataset)
    base.case_count = stats["case_count"]
    return DatasetDetailOut(**base.model_dump(), stats=stats)


@router.get("/test-match")
def test_match(
    kind: str,
    dog_id: int,
    top_k: int = 10,
    apply_breed_gate: bool = False,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> dict:
    """Score one dog against ALL dogs in the opposite category by embedding similarity (testing)."""
    if kind not in ("known", "unknown"):
        raise HTTPException(status.HTTP_400_BAD_REQUEST, "kind must be 'known' or 'unknown'")
    subject_type = SubjectType.known if kind == "known" else SubjectType.unknown
    query = ds.dog_profile(db, subject_type, dog_id)
    if query is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "Dog not found")

    top_k = max(1, min(top_k, 50))
    from ..services.matching import rank_dog_against_opposite

    ranked, considered, query_embedded = rank_dog_against_opposite(
        db, subject_type=subject_type, subject_id=dog_id, top_k=top_k,
        apply_breed_gate=apply_breed_gate,
    )
    results = []
    for cand_type, cid, score in ranked:
        prof = ds.dog_profile(db, cand_type, cid)
        if prof:
            results.append({**prof, "score": score})

    from ..ml import get_embedder

    e = get_embedder()
    return {
        "query": query,
        "results": results,
        "model": f"{e.name}/{e.version}",
        "candidate_count": considered,
        "query_embedded": query_embedded,
    }


@router.get("/dog/{kind}/{dog_id}")
def dog_detail(
    kind: str,
    dog_id: int,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> dict:
    """A dog's profile plus ALL of its photos (full images) — for the profile viewer."""
    if kind not in ("known", "unknown"):
        raise HTTPException(status.HTTP_400_BAD_REQUEST, "kind must be 'known' or 'unknown'")
    subject_type = SubjectType.known if kind == "known" else SubjectType.unknown
    profile = ds.dog_profile(db, subject_type, dog_id)
    if profile is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "Dog not found")
    from .helpers import pictures_for

    photos = [p.model_dump() for p in pictures_for(db, subject_type, dog_id)]
    return {"profile": profile, "photos": photos}


@router.get("/breeds")
def list_breeds(
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> dict:
    """Predicted breed labels (HF softmax, active model) available to filter by."""
    return {"model": ds.active_breed_model_in_db(db), "breeds": ds.list_breeds(db)}


@router.get("/dogs")
def list_all_dogs(
    kind: str = "all",
    limit: int = 30,
    offset: int = 0,
    breed: str | None = None,
    breed_k: int = 10,
    sort: str = "newest",
    zip: str | None = None,
    added_from: date | None = None,
    added_to: date | None = None,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> dict:
    """Browse every dog profile, filterable by predicted breed / ZIP / added-date and sortable.

    ``breed_k`` = 1 matches the top predicted breed; 5–10 matches within the top-K. ``sort`` is
    ``newest`` (default) or ``oldest``. ``zip`` matches a dog's ZIP by prefix; ``added_from`` /
    ``added_to`` bound the added date (inclusive).
    """
    kind = kind if kind in ("all", "known", "unknown") else "all"
    limit = max(1, min(limit, 100))
    offset = max(0, offset)
    breed_k = max(1, min(breed_k, 50))
    sort = sort if sort in ("newest", "oldest") else "newest"
    return ds.list_all_dogs(
        db, kind, limit, offset, breed=breed or None, breed_k=breed_k,
        sort=sort, zip_prefix=(zip or None), added_from=added_from, added_to=added_to,
    )


@router.get("/datasets/{dataset_id}/dogs")
def list_dataset_dogs(
    dataset_id: int,
    limit: int = 25,
    offset: int = 0,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> dict:
    dataset = _get_dataset(db, dataset_id)
    limit = max(1, min(limit, 100))
    return ds.list_dogs(db, dataset, limit, offset)


@router.delete("/datasets/{dataset_id}", response_model=PurgeResult)
def delete_dataset(
    dataset_id: int,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> PurgeResult:
    dataset = _get_dataset(db, dataset_id)
    deleted = ds.purge_dataset(db, dataset)
    return PurgeResult(dataset_id=dataset_id, deleted=deleted)


@router.get("/image/{image_id}/dataset", response_model=ImageDatasetOut)
def image_dataset(
    image_id: int,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> ImageDatasetOut:
    pic = db.get(Picture, image_id)
    if not pic:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "Image not found")
    from ..models import KnownDog, UnknownDog

    if pic.subject_type == SubjectType.known:
        dog = db.get(KnownDog, pic.subject_id)
    else:
        dog = db.get(UnknownDog, pic.subject_id)
    dataset_id = getattr(dog, "dataset_id", None) if dog else None
    dataset_name = None
    if dataset_id:
        d = db.get(Dataset, dataset_id)
        dataset_name = d.name if d else None
    return ImageDatasetOut(
        image_id=image_id,
        subject_type=pic.subject_type.value,
        subject_id=pic.subject_id,
        dataset_id=dataset_id,
        dataset_name=dataset_name,
    )


@router.post("/datasets/{dataset_id}/embed-all", response_model=EmbedAllResult)
def embed_all(
    dataset_id: int,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> EmbedAllResult:
    dataset = _get_dataset(db, dataset_id)
    result = ds.embed_all(db, dataset)
    return EmbedAllResult(dataset_id=dataset_id, **result)


@router.post("/datasets/{dataset_id}/match", response_model=MatchDatasetResult)
def match_dataset(
    dataset_id: int,
    candidate_dataset_id: int | None = None,
    _: User = Depends(require_admin),
    db: Session = Depends(get_db),
) -> MatchDatasetResult:
    dataset = _get_dataset(db, dataset_id)
    result = ds.match_dataset(db, dataset, candidate_dataset_id=candidate_dataset_id)
    return MatchDatasetResult(
        dataset_id=dataset_id,
        cases_processed=result["cases_processed"],
        matches_created=result["matches_created"],
        candidate_dataset_id=result["candidate_dataset_id"],
    )