backend/app/services/batch_loader.py

"""Batch loader: turn a folder of images + a CSV into Dataset-linked DB records.

Reuses the existing image pipeline (``process_and_store_picture``) for embedding/breed generation —
nothing is reimplemented here. One CSV row == one image == one KnownDog/UnknownDog record (matching
``prepare_test_data.py``). Owners/finders are de-duped by email and linked to the dataset.

Individual image failures are logged and skipped (the load never aborts); total errors are reported.
"""
from __future__ import annotations

import csv
import logging
import random
from dataclasses import dataclass, field
from datetime import date
from pathlib import Path
from typing import Callable

from sqlalchemy import select
from sqlalchemy.orm import Session

from ..models import Case, Dataset, KnownDog, UnknownDog, User
from ..models.base import (
    CaseStatus,
    CaseType,
    DatasetType,
    DogSize,
    KnownDogStatus,
    SubjectType,
    UnknownDogStatus,
    UserRole,
)
from .images import ImageValidationError, process_and_store_picture

logger = logging.getLogger("pawtrace.batch_loader")

ProgressCb = Callable[[dict], None]


@dataclass
class LoadResult:
    dataset_id: int
    dogs_loaded: int = 0
    images_processed: int = 0
    cases_created: int = 0
    users_created: int = 0
    errors: list[str] = field(default_factory=list)
    matching: dict | None = None

    def as_dict(self) -> dict:
        return {
            "dataset_id": self.dataset_id,
            "dogs_loaded": self.dogs_loaded,
            "images_processed": self.images_processed,
            "cases_created": self.cases_created,
            "users_created": self.users_created,
            "error_count": len(self.errors),
            "errors": self.errors,
            "matching": self.matching,
        }


def _coerce_size(value: str | None) -> DogSize | None:
    if not value:
        return None
    try:
        return DogSize(value.strip().lower())
    except ValueError:
        return None


def _read_csv(csv_path: Path) -> list[dict]:
    with csv_path.open(newline="", encoding="utf-8") as fh:
        return list(csv.DictReader(fh))


def _get_or_create_user(
    db: Session, dataset_id: int, *, name: str, email: str, phone: str | None, zip_code: str,
    role: UserRole, created_counter: list[int],
) -> User:
    email = (email or "").strip().lower()
    user = None
    if email:
        user = db.execute(select(User).where(User.email == email)).scalar_one_or_none()
    if user is None:
        # Synthesize a unique email if the CSV omitted one (anonymous finder).
        if not email:
            email = f"anon-{dataset_id}-{created_counter[0]}@example.invalid"
        user = User(
            name=name or "Unknown", email=email, phone=phone or None,
            zip=zip_code or "00000", role=role, dataset_id=dataset_id,
        )
        db.add(user)
        db.flush()
        created_counter[0] += 1
    return user


def _resolve_image(root: Path, folder: str, image_file: str) -> Path | None:
    candidate = root / folder / image_file
    if candidate.is_file():
        return candidate
    return None


def load_dataset(
    db: Session,
    *,
    folder: str | Path,
    dataset_type: str,
    name: str,
    description: str | None,
    csv_path: str | Path,
    mark_lost: bool = False,
    mark_lost_pct: int = 100,
    run_matching: bool = False,
    skip_embeddings: bool = False,
    group_by_folder: bool = True,
    seed: int = 42,
    progress: ProgressCb | None = None,
) -> LoadResult:
    root = Path(folder)
    csv_file = Path(csv_path)
    if not csv_file.is_file():
        raise FileNotFoundError(f"CSV not found: {csv_file}")

    dtype = DatasetType(dataset_type)
    is_known = dtype in (DatasetType.known, DatasetType.test_known)

    dataset = Dataset(
        name=name, type=dtype, description=description, source_path=str(root), dog_count=0,
    )
    db.add(dataset)
    db.flush()

    rows = _read_csv(csv_file)
    total = len(rows)
    rng = random.Random(seed)
    result = LoadResult(dataset_id=dataset.id)
    created_counter = [0]

    # Group rows so every image of one identity becomes ONE dog with multiple pictures. The CSV
    # `folder` column is the identity key (each DogFaceNet subfolder == one dog). With
    # group_by_folder=False each row becomes its own single-picture dog (legacy behavior).
    if group_by_folder:
        grouped: dict[str, list[dict]] = {}
        for row in rows:
            grouped.setdefault((row.get("folder") or "").strip(), []).append(row)
        groups = list(grouped.items())
    else:
        groups = [((row.get("folder") or "").strip(), [row]) for row in rows]

    done_rows = 0

    def _emit() -> None:
        if progress:
            progress({"processed": done_rows, "total": total, "dataset_id": dataset.id})

    for folder_name, frows in groups:
        try:
            if is_known:
                _load_known_folder(db, dataset, folder_name, frows, root, rng,
                                   mark_lost, mark_lost_pct, result, created_counter,
                                   generate=not skip_embeddings)
            else:
                _load_unknown_folder(db, dataset, folder_name, frows, root, result,
                                     created_counter, generate=not skip_embeddings)
        except Exception as exc:  # noqa: BLE001 - skip-and-continue per spec
            msg = f"folder {folder_name!r}: {exc}"
            logger.warning(msg, exc_info=True)
            result.errors.append(msg)
        done_rows += len(frows)
        _emit()

    dataset.dog_count = result.dogs_loaded
    db.commit()

    if run_matching:
        from .datasets import match_dataset

        result.matching = match_dataset(db, dataset)

    return result


def _load_known_folder(db, dataset, folder_name, frows, root, rng, mark_lost, mark_lost_pct,
                       result, created_counter, *, generate: bool = True) -> None:
    """Create one KnownDog for the folder, attaching every readable image as a picture."""
    meta = frows[0]
    dog: KnownDog | None = None
    owner: User | None = None
    first_pic = True
    for row in frows:
        image_file = (row.get("image_file") or "").strip()
        img_path = _resolve_image(root, folder_name, image_file)
        if img_path is None:
            result.errors.append(f"folder {folder_name!r}: image not found {image_file!r}")
            continue
        try:
            data = img_path.read_bytes()
            if dog is None:
                owner = _get_or_create_user(
                    db, dataset.id,
                    name=meta.get("owner_name", ""), email=meta.get("owner_email", ""),
                    phone=meta.get("owner_phone"), zip_code=meta.get("zip", ""),
                    role=UserRole.owner, created_counter=created_counter,
                )
                dog = KnownDog(
                    owner_id=owner.id, dataset_id=dataset.id,
                    name=meta.get("dog_name") or "Unnamed",
                    breed=meta.get("breed") or None,
                    color=meta.get("color") or None,
                    size=_coerce_size(meta.get("size")),
                    description=meta.get("description") or "",
                    status=KnownDogStatus.home,
                )
                db.add(dog)
                db.flush()
                result.dogs_loaded += 1
                result.users_created = created_counter[0]
            process_and_store_picture(
                db, subject_type=SubjectType.known, subject_id=dog.id, data=data,
                is_primary=first_pic, generate_embedding=generate, generate_breed=generate,
            )
            first_pic = False
            result.images_processed += 1
        except ImageValidationError as exc:
            result.errors.append(f"folder {folder_name!r} image {image_file!r}: invalid ({exc})")
        except Exception as exc:  # noqa: BLE001
            result.errors.append(f"folder {folder_name!r} image {image_file!r}: {exc}")

    if dog is not None and mark_lost and rng.random() < (mark_lost_pct / 100.0):
        zip_code = meta.get("zip") or (owner.zip if owner else "")
        dog.status = KnownDogStatus.lost
        dog.last_known_zip = zip_code
        db.add(
            Case(
                person_id=owner.id if owner else None, known_dog_id=dog.id, type=CaseType.lost,
                event_zip=zip_code, event_date=date.today(),
                search_radius_miles=0, status=CaseStatus.open,
            )
        )
        result.cases_created += 1


def _load_unknown_folder(db, dataset, folder_name, frows, root, result, created_counter,
                         *, generate: bool = True) -> None:
    """Create one UnknownDog + one found Case for the folder, attaching every readable image."""
    meta = frows[0]
    dog: UnknownDog | None = None
    finder: User | None = None
    first_pic = True
    for row in frows:
        image_file = (row.get("image_file") or "").strip()
        img_path = _resolve_image(root, folder_name, image_file)
        if img_path is None:
            result.errors.append(f"folder {folder_name!r}: image not found {image_file!r}")
            continue
        try:
            data = img_path.read_bytes()
            if dog is None:
                finder = _get_or_create_user(
                    db, dataset.id,
                    name=meta.get("finder_name", ""), email=meta.get("finder_email", ""),
                    phone=meta.get("finder_phone"), zip_code=meta.get("found_zip", ""),
                    role=UserRole.finder, created_counter=created_counter,
                )
                dog = UnknownDog(
                    dataset_id=dataset.id,
                    description=meta.get("description") or "",
                    color=meta.get("color") or None,
                    size=_coerce_size(meta.get("size")),
                    current_zip=meta.get("found_zip") or "",
                    current_location_detail=meta.get("current_location") or None,
                    status=UnknownDogStatus.at_shelter,
                )
                db.add(dog)
                db.flush()
                result.dogs_loaded += 1
                result.users_created = created_counter[0]
            process_and_store_picture(
                db, subject_type=SubjectType.unknown, subject_id=dog.id, data=data,
                is_primary=first_pic, generate_embedding=generate, generate_breed=generate,
            )
            first_pic = False
            result.images_processed += 1
        except ImageValidationError as exc:
            result.errors.append(f"folder {folder_name!r} image {image_file!r}: invalid ({exc})")
        except Exception as exc:  # noqa: BLE001
            result.errors.append(f"folder {folder_name!r} image {image_file!r}: {exc}")

    if dog is not None and finder is not None:
        found_zip = meta.get("found_zip") or ""
        db.add(
            Case(
                person_id=finder.id,
                finder_name=meta.get("finder_name") or finder.name,
                finder_email=meta.get("finder_email") or finder.email,
                finder_phone=meta.get("finder_phone"),
                unknown_dog_id=dog.id, type=CaseType.found,
                event_zip=found_zip, event_date=date.today(),
                current_location=meta.get("current_location") or None,
                search_radius_miles=0, status=CaseStatus.open,
            )
        )
        result.cases_created += 1