photoanalyzer/photo_pipeline/services/safety.py

"""Safety (NSFW) logic: scoring bands, decisions, and EXIF keyword rules.

Pure, deterministic functions extracted from the donors so the model adapter and
exiftool adapter stay at the edges. This is the shared source of truth for how a
score becomes a decision and how safety keywords are read and projected.

Extracted from (donor_ledger.yaml):
- pa-nsfw-filter — photo_analyzer._rec_has_nsfw / filter_nsfw_tagged
- nt-exif-marks  — nsfwtag.exif.read_marks / read_tagged (keyword parsing)
- nsfwtag thresholds (DEFAULT_THRESHOLD / DEFAULT_REVIEW_MIN)

Behavior changes vs the donors (required by stable identity + shared jobs):
- Decisions and scores key on ``asset_id`` and live in the database, not in a
  path-keyed ``nsfw_scores.csv`` (the CSV is a one-time migration import).
- Safety keywords are **mutually exclusive**: writing a decision adds the chosen
  ``sfw``/``nsfw`` keyword and removes the opposite. The donor only ever added
  ``nsfw``; there was no ``sfw`` keyword and no removal.
- Model scoring runs behind ``integrations.nsfw_model.NsfwModel`` as a shared job,
  not inline in a CLI pass.
- No module here imports the archived CLIs (``nsfwtag`` / ``photo_analyzer``).
"""

from __future__ import annotations

from collections.abc import Iterable, Mapping

NSFW = "nsfw"
SFW = "sfw"
REVIEW_REQUIRED = "review_required"

DEFAULT_THRESHOLD = 0.6  # AdamCodd model sweet spot (donor DEFAULT_THRESHOLD)
DEFAULT_REVIEW_MIN = 0.2  # below this is confidently SFW (donor DEFAULT_REVIEW_MIN)


def classify(
    score: float, *, threshold: float = DEFAULT_THRESHOLD, review_min: float = DEFAULT_REVIEW_MIN
) -> str:
    """Suggested decision from a model score. A human confirms before it is durable.

    ``>= threshold`` → ``nsfw``; ``>= review_min`` → ``review_required``; else ``sfw``.
    """
    if score >= threshold:
        return NSFW
    if score >= review_min:
        return REVIEW_REQUIRED
    return SFW


def normalize_keywords(values: Iterable | str | None) -> set[str]:
    """Normalize an EXIF Keywords/Subject value (list or scalar) to a lowercased set."""
    if values is None:
        return set()
    if isinstance(values, str):
        values = [values]
    return {str(v).strip().lower() for v in values}


def has_nsfw(keywords: Iterable[str]) -> bool:
    """True if the ``nsfw`` safety keyword is present (donor _rec_has_nsfw)."""
    return NSFW in {str(k).strip().lower() for k in keywords}


def marks_from_keywords(keyword_sets: Mapping[str, Iterable[str]]) -> dict[str, set[str]]:
    """Split ``{path: keywords}`` into ``{'nsfw': set, 'sfw': set}``.

    They should be mutually exclusive; if a file carries both, ``nsfw`` wins so it
    never reads as safe (donor nsfwtag.exif.read_marks).
    """
    nsfw, sfw = set(), set()
    for path, keywords in keyword_sets.items():
        lowered = {str(k).strip().lower() for k in keywords}
        if NSFW in lowered:
            nsfw.add(path)
        elif SFW in lowered:
            sfw.add(path)
    return {NSFW: nsfw, SFW: sfw}


def partition_nsfw(
    keyword_sets: Mapping[str, Iterable[str]],
) -> tuple[list[str], list[str]]:
    """Split paths into ``(analyzable, skipped)`` — nsfw-tagged files skip cloud
    analysis (privacy gate; donor photo_analyzer.filter_nsfw_tagged)."""
    analyzable, skipped = [], []
    for path, keywords in keyword_sets.items():
        (skipped if has_nsfw(keywords) else analyzable).append(path)
    return analyzable, skipped


def exif_projection(decision: str) -> dict[str, list[str]]:
    """Keyword operations to make EXIF match a confirmed decision.

    Enforces mutual exclusion: add the decision keyword, remove the opposite.
    Non-terminal decisions (review/unknown) touch nothing.
    """
    if decision == NSFW:
        return {"add": [NSFW], "remove": [SFW]}
    if decision == SFW:
        return {"add": [SFW], "remove": [NSFW]}
    return {"add": [], "remove": []}


# ── Persistence, scoring, and the durable safety decision ────────────────────

import uuid
from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.orm import sessionmaker

from photo_pipeline.integrations import exiftool
from photo_pipeline.models import Asset, SafetyReview
from photo_pipeline.services import hashing

DECISIONS = {SFW, NSFW, "deferred"}


class SafetyError(Exception):
    """Invalid safety operation (bad decision, unknown asset)."""


def _now() -> datetime:
    # Microsecond precision so the latest-review-wins ordering never ties, even
    # when a decision is revised twice in the same DB-second (SQLite CURRENT_TIMESTAMP
    # is only second-granular).
    return datetime.now(timezone.utc)


class SafetyService:
    """Durable safety scoring, review decisions, and the EXIF safety checkpoint.

    Scores and decisions live in ``safety_reviews`` keyed by ``asset_id`` (the
    donor's path-keyed CSV is gone). The latest row per asset is the current
    decision; ``AnalysisService`` reads it as the privacy gate.
    """

    def __init__(self, session_factory: sessionmaker, *, model=None) -> None:
        self._session_factory = session_factory
        self._model = model  # injected in tests; real NsfwModel is lazy-loaded

    # -- reads ----------------------------------------------------------------
    def _latest_by_asset(self, session) -> dict[str, SafetyReview]:
        # Latest row per asset. Small local scale: order ascending, let later rows
        # overwrite. ponytail: a windowed query if safety_reviews ever grows huge.
        latest: dict[str, SafetyReview] = {}
        for review in session.scalars(select(SafetyReview).order_by(SafetyReview.created_at)):
            latest[review.asset_id] = review
        return latest

    def current_decision(self, asset_id: str) -> str | None:
        with self._session_factory() as session:
            review = self._latest_by_asset(session).get(asset_id)
            return review.decision if review else None

    def counts(self) -> dict[str, int]:
        """Decision breakdown over canonical, active assets — the workflow totals."""
        with self._session_factory() as session:
            assets = list(session.scalars(_eligible_assets_query()))
            latest = self._latest_by_asset(session)
            out = {SFW: 0, NSFW: 0, "deferred": 0, "undecided": 0, "scored": 0}
            for asset in assets:
                review = latest.get(asset.id)
                decision = review.decision if review else None
                if decision in (SFW, NSFW, "deferred"):
                    out[decision] += 1
                else:
                    out["undecided"] += 1
                    if review and review.score is not None:
                        out["scored"] += 1
            return out

    def review_queue(self, state: str = "", limit: int = 100, offset: int = 0) -> dict:
        """Assets for the review UI, filtered by ``state`` (undecided/sfw/nsfw/deferred)."""
        with self._session_factory() as session:
            assets = list(session.scalars(_eligible_assets_query().order_by(Asset.current_path)))
            latest = self._latest_by_asset(session)
            rows = []
            for asset in assets:
                review = latest.get(asset.id)
                decision = review.decision if review else None
                effective = decision or "undecided"
                if state and state != effective:
                    continue
                rows.append(
                    {
                        "asset_id": asset.id,
                        "current_path": asset.current_path,
                        "score": review.score if review else None,
                        "decision": decision,
                        "suggested": classify(review.score) if review and review.score is not None else None,
                        "exif_verified": bool(review and review.exif_verified_at),
                    }
                )
            return {"total": len(rows), "items": rows[offset : offset + limit]}

    def scorable_asset_ids(self) -> list[str]:
        """Canonical active assets with a path — the items a scoring job enqueues."""
        with self._session_factory() as session:
            return [a.id for a in session.scalars(_eligible_assets_query()) if a.current_path]

    # -- writes ---------------------------------------------------------------
    def score_assets(self, asset_ids: list[str] | None = None) -> int:
        """Score eligible assets with the local model and persist scored reviews.

        Records score-only reviews (no decision) so the reviewer can sort by risk.
        Returns the number scored. Assets already carrying a human decision keep it.
        """
        with self._session_factory() as session:
            query = _eligible_assets_query()
            if asset_ids is not None:
                query = query.where(Asset.id.in_(asset_ids))
            assets = [a for a in session.scalars(query) if a.current_path]
        if not assets:
            return 0
        model = self._model or _default_model()
        by_path = {a.current_path: a.id for a in assets}
        scored = model.score(list(by_path))
        now = _now()
        with self._session_factory() as session:
            for path, score in scored:
                session.add(
                    SafetyReview(
                        id=str(uuid.uuid4()), asset_id=by_path[path], score=float(score), created_at=now
                    )
                )
            session.commit()
        return len(scored)

    def decide(self, asset_id: str, decision: str, *, reviewer: str = "user", write_exif: bool = True) -> dict:
        """Persist a human safety decision and, for sfw/nsfw, run the EXIF checkpoint.

        The decision is durable regardless of EXIF; ``exif_verified`` becomes true
        only after the mutually-exclusive keyword is written and read back, and the
        file's new SHA-256 is stored — that verified state is what upload eligibility
        later depends on.
        """
        if decision not in DECISIONS:
            raise SafetyError(f"invalid decision {decision!r}")
        with self._session_factory() as session:
            asset = session.get(Asset, asset_id)
            if asset is None:
                raise SafetyError(f"unknown asset {asset_id}")
            prior = self._latest_by_asset(session).get(asset_id)
            path = asset.current_path

        exif_verified_at = None
        result_sha256 = None
        if write_exif and decision in (SFW, NSFW) and path:
            ops = exif_projection(decision)
            if exiftool.apply_keywords(path, add=ops["add"], remove=ops["remove"]):
                # Read back: the chosen keyword present, the opposite absent.
                keywords = exiftool.read_keyword_sets([path]).get(path, set())
                opposite = NSFW if decision == SFW else SFW
                if decision in keywords and opposite not in keywords:
                    exif_verified_at = _now()
                    result_sha256 = hashing.sha256_file(path)

        now = _now()
        with self._session_factory() as session:
            review = SafetyReview(
                id=str(uuid.uuid4()),
                asset_id=asset_id,
                decision=decision,
                prior_decision=prior.decision if prior else None,
                reviewer=reviewer,
                score=prior.score if prior else None,
                exif_verified_at=exif_verified_at,
                result_sha256=result_sha256,
                created_at=now,
            )
            session.add(review)
            if result_sha256:
                asset = session.get(Asset, asset_id)
                asset.current_sha256 = result_sha256
            session.commit()
        return {
            "asset_id": asset_id,
            "decision": decision,
            "prior_decision": prior.decision if prior else None,
            "exif_verified": exif_verified_at is not None,
        }


def _eligible_assets_query():
    """Canonical, active assets — the safety stage runs only on these.

    ``canonical_asset_id IS NULL`` excludes non-canonical duplicate variants; the
    shared discovery service already excludes ``_IGNORE/``.
    """
    return select(Asset).where(
        Asset.canonical_asset_id.is_(None),
        Asset.availability_state == "active",
    )


def _default_model():
    from photo_pipeline.integrations.nsfw_model import NsfwModel

    return NsfwModel()