remove-ai-watermarks/tests/test_watermark_registry.py

"""Tests for the known-visible-watermark registry (localize -> fill)."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from remove_ai_watermarks import watermark_registry as reg

DOUBAO_SAMPLE = Path(__file__).resolve().parents[1] / "data" / "fixtures" / "provenance" / "doubao-1.png"


class TestCatalog:
    def test_keys(self):
        assert reg.mark_keys() == [
            "gemini",
            "doubao",
            "jimeng",
            "qwen",
            "kling",
            "yuanbao",
            "samsung",
            "runninghub",
            "baidu",
            "liblib",
            "jimeng_pill",
        ]

    def test_all_in_auto(self):
        assert all(m.in_auto for m in reg.known_marks())

    def test_marks_expose_detect_and_mask(self):
        # Every mark drives the uniform localize -> fill contract: a detect callable
        # (verdict + bbox, no mask) and a mask callable (full-frame footprint).
        for m in reg.known_marks():
            assert callable(m._detect)
            assert callable(m._mask)

    def test_locations(self):
        by_key = {m.key: m for m in reg.known_marks()}
        assert by_key["gemini"].location == "bottom-right"
        assert by_key["doubao"].location == "bottom-right"
        assert by_key["jimeng"].location == "bottom-right"
        assert by_key["yuanbao"].location == "bottom-right"
        assert by_key["samsung"].location == "bottom-left"
        assert by_key["jimeng_pill"].location == "top-left"

    def test_get_mark_unknown_raises(self):
        with pytest.raises(KeyError):
            reg.get_mark("nope")


class TestScan:
    def test_detect_marks_scans_all(self):
        img = np.zeros((256, 256, 3), np.uint8)
        keys = {d.key for d in reg.detect_marks(img)}
        assert keys == {
            "gemini",
            "doubao",
            "jimeng",
            "qwen",
            "kling",
            "yuanbao",
            "samsung",
            "runninghub",
            "baidu",
            "liblib",
            "jimeng_pill",
        }

    def test_blank_image_no_auto_mark(self):
        dets = reg.detect_marks(np.zeros((256, 256, 3), np.uint8), include_explicit=False)
        assert not any(d.detected for d in dets)

    @pytest.mark.parametrize("shape", [(1, 1, 3), (8, 8, 3), (15, 15, 3), (12, 300, 3), (300, 10, 3)])
    def test_tiny_image_no_crash(self, shape):
        """Regression: an image whose short side is < 16 px (below the Gemini template
        floor) must yield no detection, not crash. detect_marks/remove_auto_marks are
        the public visible/all/batch path; a tiny thumbnail in a batch used to take the
        whole auto pass down with an IndexError (empty candidate list dereference)."""
        img = np.full(shape, 100, np.uint8)
        assert not any(d.detected for d in reg.detect_marks(img, include_explicit=False))
        result, removed = reg.remove_auto_marks(img, backend="cv2")
        assert removed == []
        assert result.shape == img.shape

    @pytest.mark.parametrize("shape", [(0, 5), (5, 0), (0, 5, 4), (0, 0)])
    def test_forced_remove_on_empty_array_no_crash(self, shape):
        """Regression: footprint_mask ran to_bgr (cvtColor) before any size check, so a
        forced remove on a zero-size ndarray crashed (cv2.error on an empty Mat). detect
        already guarded this; footprint_mask must too. Covers the text + gemini engines."""
        empty = np.zeros(shape, np.uint8)
        for key in ("doubao", "jimeng", "qwen", "yuanbao", "samsung", "gemini"):
            _result, mask = reg.get_mark(key).remove(empty, force=True)
            assert mask is None


class TestBackendResolution:
    def test_auto_resolves_to_available_backend(self):
        assert reg.resolve_backend("auto") in ("cv2", "migan", "lama")

    def test_explicit_backend_passes_through(self):
        assert reg.resolve_backend("cv2") == "cv2"
        assert reg.resolve_backend("lama") == "lama"

    def test_cv2_fallback_warns_once(self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture):
        import logging

        from remove_ai_watermarks import region_eraser

        monkeypatch.setattr(region_eraser, "lama_available", lambda: False)
        monkeypatch.setattr(region_eraser, "migan_available", lambda: False)
        monkeypatch.setattr(reg, "_warned_cv2_fallback", False)
        with caplog.at_level(logging.WARNING):
            assert reg.preferred_inpaint_backend() == "cv2"
            assert reg.preferred_inpaint_backend() == "cv2"
        assert sum("cv2 classical inpaint" in r.message for r in caplog.records) == 1


class TestFill:
    def test_fill_erases_masked_region(self):
        # A bright square on a flat field, masked, is inpainted away (cv2 backend).
        img = np.full((128, 128, 3), 60, np.uint8)
        img[40:70, 40:70] = 240
        mask = np.zeros((128, 128), np.uint8)
        mask[36:74, 36:74] = 255
        out = reg.fill(img, mask, backend="cv2")
        assert out.shape == img.shape
        # the masked bright square is pulled toward the surrounding field
        assert int(out[55, 55].mean()) < 160

    def test_fill_empty_mask_is_noop(self):
        img = np.full((64, 64, 3), 100, np.uint8)
        out = reg.fill(img, np.zeros((64, 64), np.uint8), backend="cv2")
        assert np.array_equal(out, img)


class TestProvenanceGate:
    """The Gemini trust gate relaxes from GEMINI_SPARKLE_TRUST_CONF to
    _GEMINI_PROVENANCE_MIN_CONF when provenance confirms Google; tested
    deterministically by stubbing the engine's raw detection confidence."""

    def _stub(self, monkeypatch: pytest.MonkeyPatch, conf: float) -> None:
        from remove_ai_watermarks.gemini_engine import DetectionResult

        # `detected` mirrors the engine's own internal floor (0.35), which is
        # independent of the registry gate under test here.
        def fake_detect(image, force_size=None, *, trust_provenance=False):
            return DetectionResult(detected=conf >= 0.35, confidence=conf, region=(10, 10, 48, 48))

        monkeypatch.setattr(reg._engine("gemini"), "detect_watermark", fake_detect)

    def test_midband_conf_needs_provenance(self, monkeypatch: pytest.MonkeyPatch):
        # Comfortably inside the relaxed band: demoted without provenance, trusted with it.
        conf = (reg._GEMINI_PROVENANCE_MIN_CONF + reg.GEMINI_SPARKLE_TRUST_CONF) / 2
        self._stub(monkeypatch, conf)
        img = np.zeros((256, 256, 3), np.uint8)
        assert reg.get_mark("gemini").detect(img).detected is False
        assert reg.get_mark("gemini").detect(img, provenance=True).detected is True

    def test_below_provenance_gate_rejected_even_with_provenance(self, monkeypatch: pytest.MonkeyPatch):
        """Provenance relaxes the gate, it does not remove it.

        Guards the 2026-07-18 raise of _GEMINI_PROVENANCE_MIN_CONF (0.35 -> 0.42).
        The engine still reports `detected` down at its own 0.35 floor, so without
        the registry gate this confidence would be accepted and inpainted. Measured
        precision just below the gate was 13% (n=30, 95% CI 5-30%) on real
        Google-metadata uploads -- i.e. ~7 of 8 accepts there destroy pixels on an
        image that never carried a sparkle, and report a removal that did not happen.
        """
        # 0.38 is inside the measured 13%-precision band and above the engine's own
        # 0.35 floor, so the engine reports `detected` and only the registry gate can
        # reject it. Hardcoded on purpose: if the gate is ever lowered back under this
        # value, this test must fail on the BEHAVIOUR below, not on its own arithmetic.
        self._stub(monkeypatch, 0.38)
        img = np.zeros((256, 256, 3), np.uint8)
        assert reg.get_mark("gemini").detect(img).detected is False
        assert reg.get_mark("gemini").detect(img, provenance=True).detected is False

    def test_provenance_gate_stays_below_the_strict_gate(self):
        """The relaxed gate must actually relax, and must not collapse onto the floor."""
        assert 0.35 < reg._GEMINI_PROVENANCE_MIN_CONF < reg.GEMINI_SPARKLE_TRUST_CONF

    def test_high_conf_detected_either_way(self, monkeypatch: pytest.MonkeyPatch):
        self._stub(monkeypatch, 0.72)
        img = np.zeros((256, 256, 3), np.uint8)
        assert reg.get_mark("gemini").detect(img).detected is True
        assert reg.get_mark("gemini").detect(img, provenance=True).detected is True


@pytest.mark.skipif(not DOUBAO_SAMPLE.exists(), reason="doubao sample not present")
class TestRealSample:
    def test_doubao_sample_detected(self):
        from remove_ai_watermarks.image_io import imread

        fired = [d.key for d in reg.detect_marks(imread(DOUBAO_SAMPLE), include_explicit=False) if d.detected]
        assert "doubao" in fired

    def test_doubao_remove_returns_region(self):
        from remove_ai_watermarks.image_io import imread

        img = imread(DOUBAO_SAMPLE)
        result, region = reg.get_mark("doubao").remove(img, backend="cv2")
        assert region is not None
        assert result.shape == img.shape


class TestLocalizeFill:
    def test_clean_corner_is_untouched(self):
        # No glyph in the corner -> no mask -> remove is a no-op copy.
        img = np.zeros((512, 512, 3), np.uint8)
        result, region = reg.get_mark("doubao").remove(img, backend="cv2")
        assert region is None
        assert np.array_equal(result, img)


class TestSensitivity:
    """``resolve_trust`` turns the sensitivity policy + evidence into the per-mark
    trust level the engines consume."""

    def test_strict_never_relaxes(self):
        # even with metadata provenance, strict keeps the conservative gate
        assert (
            reg.resolve_trust("gemini", sensitivity="strict", provenance=frozenset({"gemini"}), strict_keys=set())
            == "strict"
        )

    def test_auto_relaxes_on_own_metadata(self):
        assert (
            reg.resolve_trust("gemini", sensitivity="auto", provenance=frozenset({"gemini"}), strict_keys=set())
            == "confirmed"
        )

    def test_auto_strict_without_evidence(self):
        assert reg.resolve_trust("gemini", sensitivity="auto", provenance=frozenset(), strict_keys=set()) == "strict"

    def test_auto_cross_mark_same_product(self):
        # a detected Jimeng wordmark relaxes the Jimeng pill (same product, other corner)
        assert (
            reg.resolve_trust("jimeng_pill", sensitivity="auto", provenance=frozenset(), strict_keys={"jimeng"})
            == "confirmed"
        )

    def test_auto_no_cross_mark_across_products(self):
        # a detected Jimeng wordmark must NOT relax Doubao (distinct products, same corner)
        assert (
            reg.resolve_trust("doubao", sensitivity="auto", provenance=frozenset(), strict_keys={"jimeng"}) == "strict"
        )

    def test_remove_auto_marks_accepts_all_sensitivities(self):
        blank = np.zeros((256, 256, 3), np.uint8)
        for s in ("auto", "strict"):
            _, removed = reg.remove_auto_marks(blank, sensitivity=s, backend="cv2")
            assert removed == []


class TestNoBlanketRelaxation:
    """There is NO path that relaxes a mark's gate without same-product evidence.

    ``assume_ai`` was that path and was removed 2026-07-19: it bypassed every mark's
    false-positive gate on the caller's bare assertion that the image is AI, which says
    nothing about WHICH vendor or WHERE -- exactly what the bypass is contracted to
    require. Before it carried a confidence floor it filled a phantom sparkle on 59.8%
    of genuine camera photos. A user who can SEE a mark is served by `erase --region`
    (they supply the coordinates) or `--mark <name> --no-detect` for a text mark.
    """

    def test_sensitivity_has_exactly_two_levels(self):
        import typing

        assert set(typing.get_args(reg.Sensitivity)) == {"auto", "strict"}

    def test_trust_ladder_has_no_assumed_level(self):
        import typing

        assert set(typing.get_args(reg.Trust)) == {"strict", "confirmed"}

    def test_no_sensitivity_relaxes_without_same_product_evidence(self):
        for sens in ("auto", "strict"):
            assert reg.resolve_trust("gemini", sensitivity=sens, provenance=frozenset(), strict_keys=set()) == "strict"

    def test_the_assumed_floor_helper_is_gone(self):
        """It existed only to make the blanket relaxation tolerable."""
        assert not hasattr(reg, "assumed_floor_ok")
        assert not hasattr(reg, "_ASSUMED_CONF_FLOOR")

    def test_the_removed_value_raises_instead_of_silently_meaning_auto(self):
        """`Sensitivity` is a Literal and unenforced at runtime, so a 0.15 caller passing
        the removed value would quietly get `auto` -- a silent semantic change on exactly
        the release where they need to be told. The error names the replacement."""
        import pytest

        with pytest.raises(ValueError, match="erase"):
            reg.validate_sensitivity("assume_ai")
        with pytest.raises(ValueError, match="unknown sensitivity"):
            reg.validate_sensitivity("aggressive")
        assert reg.validate_sensitivity("auto") == "auto"
        assert reg.validate_sensitivity("strict") == "strict"

    def test_context_rejects_it_too(self):
        """The arbiter's own entry point validates, so a direct `decide()` caller cannot
        smuggle the removed mode past the public API."""
        import pytest

        with pytest.raises(ValueError, match=r"removed in 0\.16"):
            reg.Context(sensitivity="assume_ai")


class TestArbiter:
    """``decide`` is the PURE removal arbiter: (candidates, context) -> ordered
    winners, no image / no I/O. Tested in isolation by handing it fabricated
    Candidates -- this is the payoff of separating decision from perception."""

    @staticmethod
    def _c(key, *, strict=False, relaxed=False, flat=False):
        feats = {"footprint_flat": 1.0} if flat else {}
        return reg.Candidate(key, f"L:{key}", strict, relaxed, feats)

    def _keys(self, cands, ctx):
        return {d.candidate.key for d in reg.decide(cands, ctx)}

    def test_empty(self):
        assert reg.decide([], reg.Context()) == []

    def test_strict_uses_strict_verdict(self):
        # relaxed-only detection must NOT fire under strict
        assert self._keys([self._c("gemini", relaxed=True)], reg.Context(sensitivity="strict")) == set()

    def test_auto_relaxes_on_provenance(self):
        c = [self._c("gemini", relaxed=True)]
        assert self._keys(c, reg.Context(provenance=frozenset({"gemini"}))) == {"gemini"}
        assert self._keys(c, reg.Context()) == set()  # no evidence -> strict verdict (not fired)

    def test_cross_mark_relaxes_pill_via_jimeng(self):
        cands = [self._c("jimeng", strict=True, relaxed=True), self._c("jimeng_pill", relaxed=True, flat=True)]
        assert self._keys(cands, reg.Context()) == {"jimeng", "jimeng_pill"}

    def test_pill_dropped_on_doubao(self):
        cands = [
            self._c("doubao", strict=True, relaxed=True),
            self._c("jimeng_pill", strict=True, relaxed=True, flat=True),
        ]
        keys = self._keys(cands, reg.Context(provenance=frozenset({"jimeng"})))
        assert "doubao" in keys
        assert "jimeng_pill" not in keys

    def test_pill_dropped_on_qwen(self):
        # A Qwen frame is TC260 too but is not Jimeng-basic either: a confident
        # bottom-right 千问AI生成 detection suppresses the pill exactly like Doubao's.
        cands = [
            self._c("qwen", strict=True, relaxed=True),
            self._c("jimeng_pill", strict=True, relaxed=True, flat=True),
        ]
        keys = self._keys(cands, reg.Context(provenance=frozenset({"jimeng"})))
        assert "qwen" in keys
        assert "jimeng_pill" not in keys

    def test_pill_dropped_on_yuanbao(self):
        # The standard Yuanbao mark identifies a different TC260 product, so a
        # coincident top-left pill match must not be treated as Jimeng-basic.
        cands = [
            self._c("yuanbao", strict=True, relaxed=True),
            self._c("jimeng_pill", strict=True, relaxed=True, flat=True),
        ]
        keys = self._keys(cands, reg.Context(provenance=frozenset({"jimeng"})))
        assert keys == {"yuanbao"}

    def test_pill_metadata_arm_gated_on_flatness(self):
        ctx = reg.Context(provenance=frozenset({"jimeng"}))
        assert self._keys([self._c("jimeng_pill", strict=True, relaxed=True, flat=True)], ctx) == {"jimeng_pill"}
        assert self._keys([self._c("jimeng_pill", strict=True, relaxed=True, flat=False)], ctx) == set()

    def test_pill_wordmark_arm_ignores_flatness(self):
        # wordmark present -> pill removed even on a textured (non-flat) footprint
        cands = [
            self._c("jimeng", strict=True, relaxed=True),
            self._c("jimeng_pill", strict=True, relaxed=True, flat=False),
        ]
        assert "jimeng_pill" in self._keys(cands, reg.Context())

    def test_weak_pill_detection_does_not_confirm_the_jimeng_wordmark(self):
        """The pill is too false-fire-prone (~7%) to grant a sibling `confirmed` trust.

        Regression: a pill false fire on clean non-ByteDance
        content confirmed jimeng, relaxing its NCC gate 0.45 -> 0.3825; jimeng then
        false-fired, and _keep_pill's wordmark arm removed the pill UNRESTRICTED,
        skipping the flatness guard. Closed loop on the default `auto` path.
        """
        # Only the pill is strictly detected. jimeng scores in the band that is
        # reachable ONLY via the relaxed gate.
        cands = [
            self._c("jimeng_pill", strict=True, relaxed=True, flat=False),
            self._c("jimeng", strict=False, relaxed=True),
        ]
        fired = {d.candidate.key for d in reg.decide(cands, reg.Context("auto", frozenset()))}
        assert "jimeng" not in fired, "a weak pill hit must not relax the jimeng wordmark"
        # and with jimeng gone, the pill loses the wordmark arm too -- no unrestricted
        # removal of a textured footprint on content nothing confirmed.
        assert "jimeng_pill" not in fired

    def test_real_jimeng_wordmark_still_corroborates_the_pill(self):
        """The fix removes only the pill's TESTIMONY, not the wordmark's."""
        cands = [
            self._c("jimeng", strict=True, relaxed=True),
            self._c("jimeng_pill", strict=True, relaxed=True, flat=False),
        ]
        fired = {d.candidate.key for d in reg.decide(cands, reg.Context("auto", frozenset()))}
        assert fired == {"jimeng", "jimeng_pill"}


class TestProvenanceMaskThreading:
    """Regression for the provenance-relaxed Gemini no-op (#1) and the false 'removed'
    label (#2). Before the fix, footprint_mask re-detected WITHOUT trust_provenance, the
    FP gate demoted the sparkle to detected=False, the mask came back None, yet
    remove_auto_marks still reported the mark as removed."""

    def test_relaxed_sparkle_yields_mask(self, monkeypatch: pytest.MonkeyPatch):
        # A sparkle a strict re-detect would demote (detected False) but a
        # provenance-relaxed detect accepts must still produce a removal mask.
        from remove_ai_watermarks.gemini_engine import DetectionResult

        def fake(image, force_size=None, *, trust_provenance=False):
            return DetectionResult(
                detected=trust_provenance, confidence=0.42 if trust_provenance else 0.30, region=(400, 400, 60)
            )

        monkeypatch.setattr(reg._engine("gemini"), "detect_watermark", fake)
        img = np.full((512, 512, 3), 90, np.uint8)
        assert reg.get_mark("gemini").localize(img, provenance=True).mask is not None
        assert reg.get_mark("gemini").localize(img, provenance=False).mask is None

    def test_no_label_when_mask_none(self, monkeypatch: pytest.MonkeyPatch):
        # A decided mark whose mask comes back None must NOT be reported as removed.
        from remove_ai_watermarks.gemini_engine import DetectionResult

        eng = reg._engine("gemini")
        monkeypatch.setattr(
            eng,
            "detect_watermark",
            lambda image, force_size=None, *, trust_provenance=False: DetectionResult(True, 0.9, (10, 10, 40)),
        )
        monkeypatch.setattr(eng, "footprint_mask", lambda image, *, force=False, region=None, dilate=None: None)
        _, removed = reg.remove_auto_marks(np.zeros((256, 256, 3), np.uint8), sensitivity="strict", backend="cv2")
        assert "Google Gemini sparkle" not in removed