tests/test_restore.py

"""Offline tests for app.restore (physics-anchored generative restoration).

All network is stubbed via the ``restore_fn`` hook; every assert compares two
independently computed quantities or checks a soft-fail contract.
"""

import numpy as np
import pytest

from app.restore import restore_layer, recover, _SEPARATE, _RESTORE, RecoverResult


def _toy_scan(seed: int = 0):
    """A tiny synthetic positive with densitometry, via the real preprocessing path."""
    from PIL import Image
    from app.preprocessing import preprocess_negative

    rng = np.random.default_rng(seed)
    # two overlapping soft blobs so a heuristic split has something to bite on
    yy, xx = np.mgrid[0:96, 0:96]
    a = np.exp(-((xx - 30) ** 2 + (yy - 40) ** 2) / 500.0)
    b = np.exp(-((xx - 65) ** 2 + (yy - 55) ** 2) / 700.0)
    img = np.clip(0.25 + 0.6 * a + 0.5 * b, 0, 1).astype(np.float32)
    rgb = np.stack([img, img * 0.9, img * 0.8], axis=-1)
    pil = Image.fromarray((rgb * 255).astype(np.uint8))
    return preprocess_negative(pil, stock="Portra 400", scan_type="positive",
                               scan_calibration="auto_exposed", auto_trim=False)


def test_restore_layer_soft_fails_without_key(monkeypatch):
    """No REPLICATE_API_TOKEN and no hook -> (None, meta) with api_contacted False."""
    monkeypatch.delenv("REPLICATE_API_TOKEN", raising=False)
    notes: list[str] = []
    out, meta = restore_layer(np.zeros((8, 8, 3), np.float32), "x", notes=notes)
    assert out is None
    assert meta["api_contacted"] is False
    assert any("REPLICATE_API_TOKEN" in n for n in notes)


def test_restore_layer_prompt_mode_selects_template():
    """separate vs restore choose different instructions, both naming the scene."""
    seen = {}

    def hook(rgb, prompt):
        seen["prompt"] = prompt
        return rgb

    _out, _m = restore_layer(np.zeros((4, 4, 3), np.float32), "a red barn",
                             mode="separate", restore_fn=hook)
    assert seen["prompt"].startswith(_SEPARATE.split("{")[0])
    assert "a red barn" in seen["prompt"]

    _out, _m = restore_layer(np.zeros((4, 4, 3), np.float32), "a red barn",
                             mode="restore", restore_fn=hook)
    assert seen["prompt"].startswith(_RESTORE.split("{")[0])
    # The two modes must not produce the same instruction (teeth).
    assert _SEPARATE.split("{")[0] != _RESTORE.split("{")[0]


def test_recover_orchestration_with_hook():
    """Hook path returns both layers, honest disclosure, and never contacts the API."""
    pre = _toy_scan()

    calls = {"n": 0}

    def hook(rgb, prompt):
        calls["n"] += 1
        return np.clip(rgb * 1.2, 0, 1)  # a distinct, non-identity 'restore'

    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, restore_fn=hook)
    assert isinstance(res, RecoverResult)
    assert res.dominant is not None and res.second is not None
    assert res.api_contacted is False
    assert calls["n"] == 2  # exactly one call per layer
    # dreamed_frac is a real fraction independently recomputable from the mask subtraction
    from app.asymmetric import subtract_layer
    _hb, _g, snr = subtract_layer(res.anchor_rgb, pre.h_total, pre.confidence_mask)
    assert res.dreamed_frac == pytest.approx(100.0 * float(np.mean(snr)), abs=1e-6)


def test_recover_without_densitometry_skips_second_layer():
    """No h_total -> dominant still attempted, second is None, note explains."""
    pre = _toy_scan()
    notes: list[str] = []
    res = recover(pre.rgb, None, None, restore_fn=lambda rgb, p: rgb, notes=notes)
    assert res.dominant is not None
    assert res.second is None
    assert any("Second layer unavailable" in n for n in notes)


def test_recover_soft_fails_offline(monkeypatch):
    """No key, no hook -> both layers None, nothing sent."""
    monkeypatch.delenv("REPLICATE_API_TOKEN", raising=False)
    pre = _toy_scan()
    res = recover(pre.rgb, pre.h_total, pre.confidence_mask)
    assert res.dominant is None and res.second is None
    assert res.api_contacted is False


def test_recover_user_scenes_win_over_vlm():
    """User-supplied scene descriptions are used verbatim and the VLM is never called."""
    pre = _toy_scan()
    prompts: list[str] = []

    def hook(rgb, prompt):
        prompts.append(prompt)
        return rgb

    def vlm_spy(pil, q):  # must NOT be consulted when the user supplied scenes
        raise AssertionError("VLM called despite user-supplied scenes")

    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, restore_fn=hook,
                  vlm=vlm_spy, scene_headline="the woman in the orange top",
                  scene_other="a stone castle wall")
    assert res.dominant is not None and res.second is not None
    assert "the woman in the orange top" in prompts[0]   # headline = separate call
    assert "a stone castle wall" in prompts[1]           # residual = restore call


def test_identify_scenes_stub_vlm():
    """identify_scenes returns the two cleaned phrases from the vlm callable."""
    from app.restore import identify_scenes

    pre = _toy_scan()
    answers = iter(["a sunlit garden", "a brick house"])
    s1, s2 = identify_scenes(pre.rgb, lambda pil, q: next(answers))
    assert (s1, s2) == ("a sunlit garden", "a brick house")
    # Teeth: over-long / multi-line junk is rejected to empty (editor falls back)
    junk = iter(["x" * 500, "line1\nline2"])
    s1, s2 = identify_scenes(pre.rgb, lambda pil, q: next(junk))
    assert s1 == "" and s2 == ""


def test_restore_best_scene_handler_scene_choice(monkeypatch):
    """UI handler routes the PICKED scene into the headline prompt (Scene 2 chosen)."""
    import app.main as m
    from PIL import Image as PILImage

    pre = _toy_scan()
    bs = m._build_best_state(
        pre, type("S", (), {"image_a": pre.rgb, "image_b": pre.rgb})(), "Portra 400", 1.0, 1.0
    )
    prompts: list[str] = []

    def fake_recover(observed, h_total=None, confidence_mask=None, *, scene_headline=None,
                     scene_other=None, notes=None, **kwargs):
        from app.restore import RecoverResult
        prompts.append((scene_headline, scene_other))
        return RecoverResult(dominant=np.zeros((4, 4, 3), np.float32), notes=notes or [])

    import app.restore as r
    monkeypatch.setattr(r, "recover", fake_recover)
    dom, sec, status, _alts, _alt_state = m.restore_best_scene(bs, "Scene 2", "castle", "person")
    assert prompts == [("person", "castle")]  # Scene 2 ("person") leads
    assert dom is not None
    assert "person" in status


def test_scene_prompt_truncated_at_cap():
    """Over-long scene text hits restore_layer's SAFETY NET (MAX_TOTAL_SCENE_CHARS)
    at a word boundary, with a note — the per-piece budgets live in recover()."""
    from app.restore import restore_layer, MAX_TOTAL_SCENE_CHARS

    seen = {}
    notes: list[str] = []
    long_scene = "word " * 400  # 2000 chars, well over the safety net
    restore_layer(np.zeros((4, 4, 3), np.float32), long_scene,
                  mode="restore", restore_fn=lambda rgb, p: seen.update(p=p) or rgb,
                  notes=notes)
    # The full description does not survive; it is cut at a word boundary with an ellipsis.
    assert long_scene.strip() not in seen["p"]
    assert "…" in seen["p"]
    # Independent bound: not all 400 words survived, and the surviving run fits the net.
    assert seen["p"].count("word") < 400
    assert seen["p"].count("word") <= MAX_TOTAL_SCENE_CHARS // len("word ") + 1
    # Unlike before, the truncation is REPORTED, not silent.
    assert any("prompt budget" in n for n in notes)
    # A short description passes through untouched (no ellipsis).
    seen.clear()
    restore_layer(np.zeros((4, 4, 3), np.float32), "a green patio with two women",
                  mode="restore", restore_fn=lambda rgb, p: seen.update(p=p) or rgb)
    assert "a green patio with two women" in seen["p"] and "…" not in seen["p"]


def test_run_with_backoff_retries_throttle_only(monkeypatch):
    """429/throttled errors are retried (no real sleep); other errors raise at once."""
    from app.restore import _run_with_backoff

    monkeypatch.setattr("time.sleep", lambda s: None)
    calls = {"n": 0}

    def flaky():
        calls["n"] += 1
        if calls["n"] < 3:
            raise RuntimeError("status: 429 detail: Request was throttled.")
        return "ok"

    assert _run_with_backoff(flaky) == "ok"
    assert calls["n"] == 3

    def hard_fail():
        raise ValueError("not a throttle")

    with pytest.raises(ValueError):
        _run_with_backoff(hard_fail)

    def always_throttled():
        raise RuntimeError("throttled")

    with pytest.raises(RuntimeError):
        _run_with_backoff(always_throttled, attempts=2)


def test_drift_warning_on_repaint():
    """A hook that returns unrelated noise triggers the drift warning; a faithful
    edit (brightened copy) does not."""
    pre = _toy_scan()
    rng = np.random.default_rng(0)

    notes: list[str] = []
    recover(pre.rgb, pre.h_total, pre.confidence_mask, notes=notes,
            restore_fn=lambda rgb, p: rng.random(rgb.shape).astype(np.float32))
    assert any("drifted from your photo" in n for n in notes)

    notes2: list[str] = []
    recover(pre.rgb, pre.h_total, pre.confidence_mask, notes=notes2,
            restore_fn=lambda rgb, p: np.clip(rgb * 1.2, 0, 1))
    assert not any("drifted" in n for n in notes2)


def test_best_of_n_picks_most_faithful():
    """With n=3, the noise takes lose to the faithful take; alternates sorted desc."""
    pre = _toy_scan()
    rng = np.random.default_rng(1)
    calls = {"n": 0}

    def hook(rgb, prompt):
        calls["n"] += 1
        if calls["n"] == 2:  # only the second take respects the input
            return np.clip(rgb * 1.15, 0, 1)
        return rng.random(rgb.shape).astype(np.float32)

    res = recover(pre.rgb, None, None, restore_fn=hook, n_candidates=3)
    assert calls["n"] == 3
    assert len(res.alternates) == 3
    rs = [r for _img, r in res.alternates]
    assert rs == sorted(rs, reverse=True)
    # dominant is the faithful take: independently recompute its similarity.
    # WP-19: candidates are ranked against the physics ANCHOR, not the mixed
    # frame — r vs the mixed rewards failed separations (2026-07-19 bake-off).
    from app.restore import structural_similarity_r
    assert structural_similarity_r(res.anchor_rgb, res.dominant) == max(rs)
    assert any("Best-of-3" in n for n in res.notes)


def test_best_of_default_single_call():
    """n_candidates default (1) keeps exactly one call per layer — no cost surprise."""
    pre = _toy_scan()
    calls = {"n": 0}

    def hook(rgb, prompt):
        calls["n"] += 1
        return rgb

    recover(pre.rgb, pre.h_total, pre.confidence_mask, restore_fn=hook)
    assert calls["n"] == 2  # headline + second layer, one each


def test_promote_alternate_returns_clicked_take():
    """Clicking take N in the gallery returns that take's image (pairs contract)."""
    import app.main as m

    takes = [(np.full((4, 4, 3), v, np.float32), 0.9 - 0.1 * k)
             for k, v in enumerate((0.2, 0.5, 0.8))]
    ev = type("E", (), {"index": 2})()
    img, _status = m.promote_alternate(ev, takes, "")
    assert img is not None
    assert np.allclose(np.asarray(img, np.float32) / 255.0, 0.8, atol=0.01)
    # out-of-range and empty stay no-ops
    ev_bad = type("E", (), {"index": 9})()
    upd, upd2 = m.promote_alternate(ev_bad, takes, "")
    assert not isinstance(upd, type(img))
    upd3, upd4 = m.promote_alternate(ev, None, "")
    assert not isinstance(upd3, type(img))


def test_budget_composition_keeps_hints_and_context_with_long_description():
    """A max-length scene description must not crowd out tag hints or photo context
    (the 'my text got cut off' regression): all three pieces reach the prompt."""
    from app.restore import MAX_SCENE_CHARS

    pre = _toy_scan()
    h, w = pre.rgb.shape[:2]
    sa = np.zeros((h, w), bool); sa[10:20, 10:30] = True
    sb = np.zeros((h, w), bool); sb[60:70, 60:80] = True
    long_desc = ("green table patio detail " * 40).strip()  # ~1000 chars, > cap
    prompts: list[str] = []
    notes: list[str] = []
    recover(pre.rgb, pre.h_total, pre.confidence_mask,
            restore_fn=lambda rgb, p: prompts.append(p) or rgb,
            scene_headline=long_desc, scene_other="museum wall",
            context="honeymoon roll, patio over museum",
            seeds_headline=sa, seeds_other=sb,
            hints_headline="the pool (top left)", hints_other="the frames (top right)",
            notes=notes)
    assert len(prompts) == 2
    # hints AND context both survived the over-long description
    assert "the pool (top left)" in prompts[0]
    assert "honeymoon roll" in prompts[0]
    # the description itself was capped with a user-visible note
    assert any("longer than" in n and "description" in n for n in notes)
    # and no silent tail-slice fired in restore_layer
    assert not any("prompt budget" in n for n in notes)


def test_canvas_from_upload_matches_preprocess_geometry_for_huge_upload(monkeypatch):
    """The marking canvas must share geometry with the working image even when the
    intake megapixel guard downscales a huge upload."""
    import warnings
    from PIL import Image as PILImage
    import app.preprocessing as pp
    from app.preprocessing import preprocess_negative

    # Shrink the guard so the test doesn't need a real >50MP allocation
    monkeypatch.setattr(pp, "INTAKE_MAX_MEGAPIXELS", 0.02)  # 20k px cap
    big = PILImage.fromarray(
        (np.random.default_rng(0).random((300, 400, 3)) * 255).astype(np.uint8)
    )  # 0.12 MP > cap
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        pre = preprocess_negative(big, stock="Portra 400", scan_type="positive")
        # the canvas/tap/repair surfaces all share this module-level intake
        import app.main as m
        canvas = m._working_frame_from_upload(big)
    assert canvas is not None
    assert canvas.shape[:2] == pre.rgb.shape[:2], (
        f"canvas {canvas.shape[:2]} != working image {pre.rgb.shape[:2]} — "
        f"scribble strokes would be misregistered"
    )


def test_vlm_gated_on_replicate_token_p0_privacy(monkeypatch):
    """P0-privacy: with ANTHROPIC available but NO Replicate token, the photo must
    NOT be sent to Anthropic (its descriptions could never be consumed)."""
    monkeypatch.delenv("REPLICATE_API_TOKEN", raising=False)
    pre = _toy_scan()
    calls = {"n": 0}

    def vlm_spy(pil, q):
        calls["n"] += 1
        return "a scene"

    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, vlm=vlm_spy)
    assert calls["n"] == 0, "photo sent to Anthropic despite no Replicate token"
    assert res.api_contacted is False


def test_vlm_receives_context_and_fills_only_missing(monkeypatch):
    """VLM fallback passes the user's context through and only fills MISSING scenes."""
    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")
    pre = _toy_scan()
    seen_questions: list[str] = []

    def vlm_spy(pil, q):
        seen_questions.append(q)
        return "vlm-scene"

    prompts: list[str] = []
    # restore_fn present would skip the VLM; instead patch the network boundary.
    import app.restore as r
    monkeypatch.setattr(
        r, "_replicate_restore",
        lambda rgb, prompt, slug, max_side, meta=None, refs=None: (prompts.append(prompt), rgb)[1],
    )
    notes: list[str] = []
    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, vlm=vlm_spy,
                  scene_headline="user headline", scene_other=None,
                  context="honeymoon roll context", notes=notes)
    # context reached the VLM call (identify_scenes folds it into the question)
    assert any("honeymoon roll context" in q for q in seen_questions)
    # supplied description survives; only the missing one came from the VLM
    assert any("user headline" in p for p in prompts)
    assert any("vlm-scene" in p for p in prompts)
    # Anthropic contact is disclosed
    assert res.api_contacted is True
    assert any("Anthropic" in n for n in notes)


# ---------------------------------------------------------------------------
# WP-18 D2 — status honesty
# ---------------------------------------------------------------------------

def test_promote_alternate_rerenders_drift_line(monkeypatch):
    """D2a: promoting a take rewrites the status for THAT take; a drifted take
    gains the warning; list-style gallery indices are normalized."""
    import app.main as m
    from app.restore import DRIFT_R_THRESHOLD

    good = np.full((4, 4, 3), 0.5, np.float32)
    bad = np.full((4, 4, 3), 0.9, np.float32)
    alts = [(good, 0.85), (bad, DRIFT_R_THRESHOLD - 0.1)]
    ev_ok = type("E", (), {"index": 1})()
    img, status = m.promote_alternate(ev_ok, alts, "old text\n\n**Showing take 1** (faithfulness r=0.85)")
    assert img is not None
    assert "Showing take 2" in status and "drifted" in status
    assert status.count("**Showing take") == 1  # old line stripped, one current line
    # list-style index (row, col)
    ev_list = type("E", (), {"index": [0, 0]})()
    img2, status2 = m.promote_alternate(ev_list, alts, "x")
    assert img2 is not None and "Showing take 1" in status2 and "drifted" not in status2


def test_second_scene_copy_names_densitometry_failure(monkeypatch):
    """D2b: with h_total=None the status blames densitometry, not rate limits."""
    import app.main as m

    pre = _toy_scan()
    bs = {"observed_rgb": pre.rgb, "h_total": None, "confidence_mask": None,
          "film_stock": "Portra 400", "trim_bbox_frac": None}
    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")
    import app.restore as r
    monkeypatch.setattr(r, "_replicate_restore",
                        lambda rgb, prompt, slug, max_side, meta=None: rgb)
    _dom, _sec, status, _g, _s = m.restore_best_scene(bs, "Scene 1", "a", "b")
    assert "densitometry" in status
    assert "rate-limit" not in status


def test_api_contacted_false_on_prenetwork_failure(monkeypatch):
    """D2c: an exception BEFORE the network call must not claim the photo was sent."""
    import app.restore as r

    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")

    def boom(rgb, prompt, slug, max_side, meta=None):
        raise RuntimeError("pre-network failure (no upload happened)")

    monkeypatch.setattr(r, "_replicate_restore", boom)
    notes: list[str] = []
    out, meta = r.restore_layer(np.zeros((4, 4, 3), np.float32), "x", notes=notes)
    assert out is None
    assert meta["api_contacted"] is False  # nothing left the process
    assert any("error" in n.lower() for n in notes)


# ---------------------------------------------------------------------------
# WP-18 D3 — robustness
# ---------------------------------------------------------------------------

def test_backoff_classifies_throttle_by_status_or_word(monkeypatch):
    """D3a: structured 429 retried; 'throttled' text retried; a bare '429'
    substring in an id/URL raises immediately (no 26s of sleeps)."""
    from app.restore import _run_with_backoff

    sleeps: list[float] = []
    monkeypatch.setattr("time.sleep", lambda s: sleeps.append(s))

    class Structured(Exception):
        status = 429

    calls = {"n": 0}
    def structured_flaky():
        calls["n"] += 1
        if calls["n"] < 2:
            raise Structured("rate limited")
        return "ok"
    assert _run_with_backoff(structured_flaky) == "ok"

    def id_contains_429():
        raise RuntimeError("prediction id x4295f failed: invalid input")
    sleeps.clear()
    with pytest.raises(RuntimeError):
        _run_with_backoff(id_contains_429)
    assert sleeps == []  # no retry sleeps on a hard failure


def test_slope_bounds_cached_across_fresh_curves():
    """D3b: two fresh curves of the same preset compute the LUT sweep ONCE."""
    import densitometry as d
    from film_physics import get_film_curve

    d._SLOPE_BOUNDS_CACHE.clear()
    sweeps = {"n": 0}
    orig = np.gradient

    def counting_gradient(*a, **k):
        sweeps["n"] += 1
        return orig(*a, **k)

    d.np.gradient, saved = counting_gradient, d.np.gradient
    try:
        b1 = d._slope_valid_bounds(get_film_curve("Portra 400"))
        b2 = d._slope_valid_bounds(get_film_curve("Portra 400"))  # fresh object
    finally:
        d.np.gradient = saved
    assert b1 == b2
    assert sweeps["n"] == 1, f"LUT sweep ran {sweeps['n']}x for the same preset"


def test_scene2_choice_honored_with_empty_boxes(monkeypatch):
    """D3c: empty boxes + Scene 2 picked => the VLM's SECONDARY scene leads."""
    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")
    pre = _toy_scan()
    answers = iter(["the dominant castle", "the secondary person"])
    prompts: list[str] = []
    import app.restore as r
    monkeypatch.setattr(r, "_replicate_restore",
                        lambda rgb, prompt, slug, max_side, meta=None, refs=None: (prompts.append(prompt), rgb)[1])
    r.recover(pre.rgb, pre.h_total, pre.confidence_mask,
              vlm=lambda pil, q: next(answers), headline_is_secondary=True)
    # WP-19 bundle prompts name BOTH scenes; the recover target is the one after
    # "Recover this photo:" — assert the pick leads there, not merely appears.
    assert "Recover this photo: the secondary person" in prompts[0]
    assert "Recover this photo: the dominant castle" in prompts[1]


# ---------------------------------------------------------------------------
# WP-19 — evidence-bundle restore (nano-banana-2 default)
# ---------------------------------------------------------------------------

def test_bundle_prompt_names_enemy_and_numbers_references():
    """The bundle prompt embeds the legend, numbers images to match the refs, and
    names the other scene as removable contamination (the 128 failure fix)."""
    from app.restore import build_bundle_prompt

    p = build_bundle_prompt(
        "a gallery wall with four framed prints", "a sunny patio with two women",
        has_markup=True, has_layer=True,
        legend="red strokes mark 'paintings' — belongs to this photo",
    )
    assert "Image 2 is image 1 with the user's hand-drawn annotations" in p
    assert "red strokes mark 'paintings'" in p
    assert "Image 3 is a rough physics-based separation" in p
    assert "Recover this photo: a gallery wall with four framed prints" in p
    assert "a sunny patio with two women" in p and "contamination" in p
    assert "never solidify" in p
    # without markup, the layer takes the image-2 slot
    p2 = build_bundle_prompt("a", "b", has_markup=False, has_layer=True)
    assert "Image 2 is a rough physics-based separation" in p2
    assert "Image 3" not in p2


def test_model_inputs_families():
    """kontext = single input_image; google = image_input array + 1K matched; flux-2 = input_images."""
    import io as _io
    from app.restore import _model_inputs

    bufs = [_io.BytesIO(b"a"), _io.BytesIO(b"b")]
    k = _model_inputs("black-forest-labs/flux-kontext-pro", bufs, "p")
    assert k["input_image"] is bufs[0] and "image_input" not in k

    g = _model_inputs("google/nano-banana-2", bufs, "p")
    assert g["image_input"] == bufs
    assert g["resolution"] == "1K" and g["aspect_ratio"] == "match_input_image"

    f = _model_inputs("black-forest-labs/flux-2-pro", bufs, "p")
    assert f["input_images"] == bufs and "image_input" not in f


def test_recover_threads_markup_and_layer_refs(monkeypatch):
    """On the live path the headline call carries [markup, layer] references and
    the per-scene legend reaches the prompt."""
    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")
    pre = _toy_scan()
    h, w = pre.rgb.shape[:2]
    seeds_h = np.zeros((h, w), bool); seeds_h[2:6, 2:20] = True
    seeds_o = np.zeros((h, w), bool); seeds_o[40:44, 40:60] = True
    seen: list[dict] = []

    import app.restore as r

    def spy(rgb, prompt, slug, max_side, meta=None, refs=None):
        seen.append({"prompt": prompt, "n_refs": len(refs or [])})
        return rgb

    monkeypatch.setattr(r, "_replicate_restore", spy)
    r.recover(pre.rgb, pre.h_total, pre.confidence_mask,
              scene_headline="castle", scene_other="person",
              seeds_headline=seeds_h, seeds_other=seeds_o,
              markup_rgb=pre.rgb, legend_headline="red = castle (this photo)",
              legend_other="red = castle (the other photo)")
    assert len(seen) == 2
    assert all(c["n_refs"] == 2 for c in seen)  # markup + physics layer, both calls
    assert "red = castle (this photo)" in seen[0]["prompt"]
    assert "red = castle (the other photo)" in seen[1]["prompt"]


def test_referee_ranks_candidates_over_r():
    """WP-19: with a VLM referee, the semantically best take wins even when r
    disagrees; the note discloses the referee and the Anthropic contact."""
    pre = _toy_scan()
    takes = iter([0.5, 0.7, 0.9])

    def hook(rgb, prompt):
        return np.clip(rgb * next(takes), 0, 1)

    scores = iter(["3", "9 out of 10", "5"])  # candidate 2 wins semantically
    notes: list[str] = []
    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, restore_fn=hook,
                  vlm=lambda pil, q: next(scores),
                  scene_headline="castle", scene_other="person",
                  n_candidates=3, notes=notes)
    assert len(res.alternates) == 3
    expected = np.clip(pre.rgb * 0.7, 0, 1)
    assert np.allclose(res.dominant, expected, atol=1e-6)
    assert any("referee" in n and "Anthropic" in n for n in notes)


def test_referee_failure_falls_back_to_r_ranking():
    """Every referee call failing (no digits) -> r-ranking, classic best-of note."""
    pre = _toy_scan()
    rng = np.random.default_rng(3)
    calls = {"n": 0}

    def hook(rgb, prompt):
        calls["n"] += 1
        if calls["n"] == 1:
            return np.clip(rgb * 1.1, 0, 1)  # the faithful take
        return rng.random(rgb.shape).astype(np.float32)

    notes: list[str] = []
    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, restore_fn=hook,
                  vlm=lambda pil, q: "no rating possible",
                  scene_headline="castle", scene_other="person",
                  n_candidates=3, notes=notes)
    rs = [r for _img, r in res.alternates]
    assert rs == sorted(rs, reverse=True)
    assert any("faithfulness scores" in n for n in notes)


def test_run_with_backoff_retries_transient_ssl_not_hard_errors():
    """WP-19: SSL/transport blips retry (observed live 2026-07-19); hard errors raise at once."""
    from app.restore import _run_with_backoff

    state = {"n": 0}

    def flaky():
        state["n"] += 1
        if state["n"] < 3:
            raise Exception("[SSL: SSLV3_ALERT_BAD_RECORD_MAC] ssl/tls alert bad record mac")
        return "ok"

    assert _run_with_backoff(flaky, attempts=3, wait_s=0.01) == "ok"
    assert state["n"] == 3

    def hard():
        state["n"] += 1
        raise ValueError("model does not exist")

    state["n"] = 0
    with pytest.raises(ValueError):
        _run_with_backoff(hard, attempts=3, wait_s=0.01)
    assert state["n"] == 1  # no retries burned on a hard failure


# ---------------------------------------------------------------------------
# WP-20 — identity references, corrective retry, second-scene best-of
# ---------------------------------------------------------------------------

def test_bundle_prompt_identity_role_and_numbering():
    from app.restore import build_bundle_prompt

    p = build_bundle_prompt("a", "b", has_markup=True, has_layer=True, n_identity=1)
    assert "Image 4 is a separate REAL photograph" in p
    assert "true appearance" in p and "do NOT copy its pose" in p
    p2 = build_bundle_prompt("a", "b", n_identity=1)
    assert "Image 2 is a separate REAL photograph" in p2 and "Image 3" not in p2
    # WP-21: several reference photos are numbered as a range
    p3 = build_bundle_prompt("a", "b", has_markup=True, n_identity=3)
    assert "Images 3-5 are separate REAL photographs" in p3


def test_recover_threads_identity_reference(monkeypatch):
    """identity_headline rides as the LAST reference of the headline call only."""
    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")
    pre = _toy_scan()
    h, w = pre.rgb.shape[:2]
    seeds_h = np.zeros((h, w), bool); seeds_h[2:6, 2:20] = True
    seeds_o = np.zeros((h, w), bool); seeds_o[40:44, 40:60] = True
    ident = np.full((10, 10, 3), 0.5, np.float32)
    seen: list[dict] = []

    import app.restore as r

    def spy(rgb, prompt, slug, max_side, meta=None, refs=None):
        seen.append({"prompt": prompt, "n_refs": len(refs or []),
                     "last_is_ident": bool(refs) and refs[-1].shape == (10, 10, 3)})
        return rgb

    monkeypatch.setattr(r, "_replicate_restore", spy)
    r.recover(pre.rgb, pre.h_total, pre.confidence_mask,
              scene_headline="castle", scene_other="person",
              seeds_headline=seeds_h, seeds_other=seeds_o,
              markup_rgb=pre.rgb, identity_headline=ident)
    assert seen[0]["n_refs"] == 3 and seen[0]["last_is_ident"]  # markup+layer+identity
    assert "true appearance" in seen[0]["prompt"]
    assert seen[1]["n_refs"] == 2 and not seen[1]["last_is_ident"]  # no identity_other


def test_referee_critique_drives_one_corrective_retry():
    """Best take under REFEREE_RETRY_BELOW with a critique -> exactly one retry whose
    prompt carries the critique; the retry's better score wins."""
    pre = _toy_scan()
    factors = iter([0.5, 0.6, 0.7, 0.9])
    prompts: list[str] = []

    def hook(rgb, prompt):
        prompts.append(prompt)
        return np.clip(rgb * next(factors), 0, 1)

    scores = iter(["5 | three women are visible but the scene should have two",
                   "4 | frames from the other photo remain",
                   "3 | heavy contamination",
                   "9 | none"])
    notes: list[str] = []
    res = recover(pre.rgb, None, None, restore_fn=hook,  # no h_total: headline only
                  vlm=lambda pil, q: next(scores),
                  scene_headline="castle", scene_other="person",
                  n_candidates=3, notes=notes)
    assert len(prompts) == 4  # 3 takes + 1 corrective retry
    assert "three women are visible" in prompts[3] and "Do not repeat" in prompts[3]
    assert np.allclose(res.dominant, np.clip(pre.rgb * 0.9, 0, 1), atol=1e-6)
    assert any("corrective retry" in n for n in notes)
    assert len(res.alternates) == 4  # the retry joins the gallery


def test_second_scene_joins_best_of_with_referee():
    """n_candidates=3 -> the second scene runs best-of-2, referee-ranked."""
    pre = _toy_scan()
    factors = iter([0.5, 0.7, 0.9, 0.6, 0.8])

    def hook(rgb, prompt):
        return np.clip(rgb * next(factors), 0, 1)

    scores = iter(["3 | a", "9 | none", "5 | b",   # headline: best 9, no retry
                   "8 | none", "6 | c"])           # second: first take wins
    res = recover(pre.rgb, pre.h_total, pre.confidence_mask, restore_fn=hook,
                  vlm=lambda pil, q: next(scores),
                  scene_headline="castle", scene_other="person",
                  n_candidates=3)
    assert len(res.alternates) == 3
    assert np.allclose(res.dominant, np.clip(pre.rgb * 0.7, 0, 1), atol=1e-6)
    assert res.second is not None
    assert np.allclose(res.second, np.clip(pre.rgb * 0.6, 0, 1), atol=1e-6)



# ---------------------------------------------------------------------------
# WP-21 — multi-reference identities, 2K finalize
# ---------------------------------------------------------------------------

def test_identity_list_capped_at_budget(monkeypatch):
    monkeypatch.setenv("REPLICATE_API_TOKEN", "tok")
    import app.restore as r
    pre = _toy_scan()
    seen = {}

    def spy(rgb, prompt, slug, max_side, meta=None, refs=None):
        seen.setdefault("n_refs", len(refs or []))
        return rgb

    monkeypatch.setattr(r, "_replicate_restore", spy)
    notes: list[str] = []
    ident = [np.full((8, 8, 3), 0.5, np.float32)] * 6  # over the 4-per-scene budget
    r.recover(pre.rgb, None, None, scene_headline="castle", scene_other="person",
              identity_headline=ident, notes=notes)
    assert seen["n_refs"] == 1 + 4  # anchor layer + capped identities
    assert any("reference budget" in n for n in notes)


def test_finalize_take_offline_hook_and_drift_note():
    from app.restore import finalize_take

    rgb = np.tile(np.linspace(0, 1, 64, dtype=np.float32)[None, :, None], (48, 1, 3))
    rgb[10:20, 10:30] = 0.9
    out, meta = finalize_take(rgb, restore_fn=lambda im, prompt: im * 0.98)
    assert out is not None and meta["api_contacted"] is False
    notes: list[str] = []
    rng = np.random.default_rng(7)
    finalize_take(rgb, restore_fn=lambda im, prompt: rng.random(im.shape, np.float32),
                  notes=notes)
    assert any("drifted" in n for n in notes)


def test_finalize_take_soft_fails_without_token(monkeypatch):
    from app.restore import finalize_take

    monkeypatch.delenv("REPLICATE_API_TOKEN", raising=False)
    notes: list[str] = []
    out, meta = finalize_take(np.zeros((8, 8, 3), np.float32), notes=notes)
    assert out is None and meta["api_contacted"] is False
    assert any("REPLICATE_API_TOKEN" in n for n in notes)


def test_fallback_chain_on_hard_failure(monkeypatch):
    """A rotated/404 primary editor falls through to the backups with a note."""
    import app.restore as r

    seen = []

    def fake_replicate_restore(rgb, prompt, slug, max_side, meta=None, refs=None, resolution="1K"):
        seen.append(slug)
        if "nano-banana-2" in slug:
            raise RuntimeError("404 Not Found: gemini preview retired")
        return rgb
    monkeypatch.setattr(r, "_replicate_restore", fake_replicate_restore)

    notes: list[str] = []
    meta: dict = {}
    out = r._restore_with_fallback(
        np.zeros((8, 8, 3), np.float32), "p", "google/nano-banana-2", 1024,
        meta, notes=notes,
    )
    assert out is not None
    assert seen == ["google/nano-banana-2", "google/nano-banana-pro"]
    assert meta["model"] == "google/nano-banana-pro"
    assert any("backup model" in n for n in notes)


def test_fallback_chain_reraises_throttle(monkeypatch):
    """Throttles are account-wide: no pointless fallback, the 429 surfaces."""
    import app.restore as r

    def throttled(rgb, prompt, slug, max_side, meta=None, refs=None, resolution="1K"):
        raise RuntimeError("Request was throttled (6/min)")
    monkeypatch.setattr(r, "_replicate_restore", throttled)

    import pytest
    with pytest.raises(RuntimeError, match="throttled"):
        r._restore_with_fallback(
            np.zeros((8, 8, 3), np.float32), "p", "google/nano-banana-2", 1024, {},
        )