bggpipe/tests/test_extract.py

"""Extract-stage tests. The vision callable is always a local fake — the
Anthropic API is never touched here."""

from __future__ import annotations

import base64
import io
import json

import pytest
import typer

from bggpipe.config import Config
from bggpipe.extract import (
    apply_title_edits,
    dedupe_entries,
    load_title_edits,
    load_title_splits,
    parse_vision_response,
    prepare_image,
    rebuild_artifacts,
    record_title_edit,
    record_title_split,
    replay_titles,
    run_extract,
)


def _write_image(path, size=(400, 300), fmt="JPEG", color=(200, 30, 30)):
    from PIL import Image

    img = Image.new("RGB", size, color)
    if fmt == "HEIC":
        from pillow_heif import register_heif_opener

        register_heif_opener()
        img.save(path, format="HEIF")
    else:
        img.save(path, format=fmt)
    return path


def _vision_stub(payload):
    """A VisionFn returning a canned response, recording call count."""
    calls = []

    def vision(image_b64, media_type):
        calls.append(media_type)
        return payload

    return vision, calls


# -- image preparation --------------------------------------------------


def test_prepare_image_downscales_long_edge(tmp_path):
    from PIL import Image

    photo = _write_image(tmp_path / "big.jpg", size=(4000, 3000))
    b64, media_type = prepare_image(photo)
    assert media_type == "image/jpeg"
    out = Image.open(io.BytesIO(base64.standard_b64decode(b64)))
    assert max(out.size) == 1568
    assert out.size == (1568, 1176)  # aspect ratio preserved


def test_prepare_image_keeps_small_images(tmp_path):
    from PIL import Image

    photo = _write_image(tmp_path / "small.png", size=(640, 480), fmt="PNG")
    b64, _ = prepare_image(photo)
    out = Image.open(io.BytesIO(base64.standard_b64decode(b64)))
    assert out.size == (640, 480)
    assert out.format == "JPEG"  # PNG converted to JPEG


def test_prepare_image_converts_heic(tmp_path):
    photo = _write_image(tmp_path / "iphone.heic", size=(2000, 1500), fmt="HEIC")
    b64, media_type = prepare_image(photo)
    assert media_type == "image/jpeg"
    assert len(b64) > 0


# -- defensive JSON parsing ---------------------------------------------


def test_parse_object_with_titles_and_unidentified():
    text = """```json
{"titles": [{"title_raw": "Catan", "confidence": "high"}],
 "unidentified": [{"location": "top shelf, left of Catan",
                   "partial_text": "WAR", "art_notes": "red spine"}]}
```"""
    titles, unidentified, _ = parse_vision_response(text)
    assert titles[0]["title_raw"] == "Catan"
    assert unidentified[0]["location"] == "top shelf, left of Catan"


def test_bare_array_response_shape_parses():
    text = '```json\n[{"title_raw": "Catan", "confidence": "high"}]\n```'
    titles, unidentified, _ = parse_vision_response(text)
    assert titles[0]["title_raw"] == "Catan"
    assert unidentified == []


def test_parse_tolerates_prose_around_json():
    text = 'Here are the games:\n[{"title_raw": "Wingspan"}]\nLet me know!'
    titles, _, _ = parse_vision_response(text)
    assert titles[0]["title_raw"] == "Wingspan"


def test_parse_drops_malformed_entries():
    text = (
        '{"titles": [{"title_raw": "Catan"}, {"no_title": true}, "just a string"],'
        ' "unidentified": [{}, "not a dict", {"location": "somewhere"}]}'
    )
    titles, unidentified, _ = parse_vision_response(text)
    assert len(titles) == 1
    assert len(unidentified) == 1  # empty {} and the bare string are dropped


def test_parse_raises_on_garbage():
    with pytest.raises(ValueError):
        parse_vision_response("I couldn't see any games clearly.")


# -- dedupe with edition cues -------------------------------------------


def _entry(title, photo, **cues):
    return {
        "title_raw": title,
        "confidence": cues.pop("confidence", "high"),
        "publisher_hint": cues.pop("publisher_hint", ""),
        "edition_hint": cues.pop("edition_hint", ""),
        "year_hint": cues.pop("year_hint", None),
        "language_hint": cues.pop("language_hint", ""),
        "art_notes": cues.pop("art_notes", ""),
        "source_photos": [photo],
    }


def test_dedupe_merges_same_title_compatible_cues():
    deduped = dedupe_entries(
        [
            _entry("Wingspan", "a.jpg", publisher_hint="Stonemaier Games"),
            _entry("WINGSPAN", "b.jpg"),  # no cues -> compatible
        ]
    )
    assert len(deduped) == 1
    assert deduped[0]["source_photos"] == ["a.jpg", "b.jpg"]
    assert deduped[0]["publisher_hint"] == "Stonemaier Games"  # provenance kept


def test_dedupe_keeps_conflicting_editions_separate():
    deduped = dedupe_entries(
        [
            _entry(
                "Cosmic Encounter", "a.jpg", edition_hint="42nd Anniversary Edition"
            ),
            _entry("Cosmic Encounter", "b.jpg", edition_hint="Eon 1977 edition"),
        ]
    )
    assert len(deduped) == 2  # different editions stay separate entries


def test_dedupe_conflicting_years_stay_separate():
    deduped = dedupe_entries(
        [
            _entry("Catan", "a.jpg", year_hint=1995),
            _entry("Catan", "b.jpg", year_hint=2015),
        ]
    )
    assert len(deduped) == 2


def test_dedupe_split_titles_never_merge():
    deduped = dedupe_entries(
        [_entry("Wiz-War", "a.jpg"), _entry("Wiz-War", "b.jpg")],
        split_titles={"wiz war"},
    )
    assert len(deduped) == 2  # the human said: separate physical copies


def test_edits_fix_misreads_before_dedupe():
    # a corrected misspelling merges with the correctly-read sighting
    edits = [{"match": "Hebarceos", "title_raw": "Herbaceous"}]
    deduped = dedupe_entries(
        apply_title_edits(
            [_entry("Hebarceos", "a.jpg"), _entry("Herbaceous", "b.jpg")], edits
        )
    )
    assert len(deduped) == 1
    assert deduped[0]["title_raw"] == "Herbaceous"
    assert deduped[0]["source_photos"] == ["a.jpg", "b.jpg"]


def test_edits_chain_and_target_photos():
    edits = [
        {"match": "Wiz-War", "photos": ["a.jpg"], "edition_hint": "7th Edition"},
        {"match": "Wiz-War", "title_raw": "Wiz-War!", "photos": ["a.jpg"]},
        # made later, against the renamed title — must chain onto the result
        {"match": "Wiz-War!", "photos": ["a.jpg"], "year_hint": 1997},
    ]
    entries = apply_title_edits(
        [_entry("Wiz-War", "a.jpg"), _entry("Wiz-War", "b.jpg")], edits
    )
    assert entries[0]["title_raw"] == "Wiz-War!"
    assert entries[0]["edition_hint"] == "7th Edition"
    assert entries[0]["year_hint"] == 1997
    assert entries[1] == _entry("Wiz-War", "b.jpg")  # untargeted copy untouched


def test_stores_roundtrip_and_replay_from_raw(tmp_path):
    cfg = Config(data_dir=tmp_path / "data", photos_dir=tmp_path / "photos")
    raw = cfg.extract_raw_dir
    raw.mkdir(parents=True)
    for photo in ("a.jpg", "b.jpg"):
        (raw / f"{photo}.json").write_text(
            json.dumps({"titles": [_entry("Wiz-War", photo)], "unidentified": []})
        )
    record_title_split(cfg.title_splits_path, "Wiz-War")
    record_title_split(cfg.title_splits_path, "wiz war")  # dupe, normalized away
    record_title_edit(
        cfg.title_edits_path,
        {"match": "Wiz-War", "photos": ["a.jpg"], "edition_hint": "7th Edition"},
    )
    replay_titles(cfg)
    titles = json.loads(cfg.titles_path.read_text())
    assert [e["source_photos"] for e in titles] == [["a.jpg"], ["b.jpg"]]
    assert titles[0]["edition_hint"] == "7th Edition"
    assert len(load_title_splits(cfg.title_splits_path)) == 1
    assert len(load_title_edits(cfg.title_edits_path)) == 1
    # a later full rebuild (a real extract run) honors the same stores
    rebuild_artifacts(
        raw,
        cfg.titles_path,
        cfg.unidentified_path,
        load_title_splits(cfg.title_splits_path),
        load_title_edits(cfg.title_edits_path),
    )
    assert len(json.loads(cfg.titles_path.read_text())) == 2


def test_replay_without_raw_caches_explodes_merged_entries(tmp_path):
    cfg = Config(data_dir=tmp_path / "data", photos_dir=tmp_path / "photos")
    cfg.data_dir.mkdir(parents=True)
    merged = _entry("Wiz-War", "a.jpg")
    merged["source_photos"] = ["a.jpg", "b.jpg", "c.jpg"]
    cfg.titles_path.write_text(json.dumps([merged, _entry("Catan", "a.jpg")]))
    record_title_split(cfg.title_splits_path, "Wiz-War")
    replay_titles(cfg)
    titles = json.loads(cfg.titles_path.read_text())
    by_title = {}
    for e in titles:
        by_title.setdefault(e["title_raw"], []).append(e["source_photos"])
    assert by_title["Wiz-War"] == [["a.jpg"], ["b.jpg"], ["c.jpg"]]
    assert by_title["Catan"] == [["a.jpg"]]  # non-split entries survive intact


def test_dedupe_upgrades_confidence():
    deduped = dedupe_entries(
        [
            _entry("Azul", "a.jpg", confidence="low"),
            _entry("Azul", "b.jpg", confidence="high"),
        ]
    )
    assert deduped[0]["confidence"] == "high"


# -- run_extract orchestration ------------------------------------------


def _cfg(tmp_path):
    return Config(photos_dir=tmp_path / "photos", data_dir=tmp_path / "data")


def test_run_extract_is_idempotent_and_resumable(tmp_path):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf1.jpg")
    _write_image(cfg.photos_dir / "shelf2.jpg")
    vision, calls = _vision_stub('[{"title_raw": "Catan", "confidence": "high"}]')

    result = run_extract(cfg, vision=vision)
    assert len(calls) == 2  # one vision call per photo
    assert len(result) == 1  # deduped across photos
    assert result[0]["source_photos"] == ["shelf1.jpg", "shelf2.jpg"]
    assert cfg.titles_path.exists()

    run_extract(cfg, vision=vision)
    assert len(calls) == 2  # second run: everything cached, no vision calls


def test_run_extract_only_reprocesses_one_photo(tmp_path):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf1.jpg")
    _write_image(cfg.photos_dir / "blurry.jpg")
    vision, calls = _vision_stub('[{"title_raw": "Azul"}]')
    run_extract(cfg, vision=vision)

    retake_vision, retake_calls = _vision_stub(
        '[{"title_raw": "Azul: Summer Pavilion"}]'
    )
    result = run_extract(cfg, only="blurry.jpg", vision=retake_vision)
    assert retake_calls == ["image/jpeg"]  # exactly one re-extraction
    titles = {e["title_raw"] for e in result}
    assert titles == {"Azul", "Azul: Summer Pavilion"}


def test_run_extract_empty_photos_dir_exits(tmp_path):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    vision, _ = _vision_stub("[]")
    with pytest.raises(typer.Exit):
        run_extract(cfg, vision=vision)


# -- unidentified sightings ---------------------------------------------


OBJECT_PAYLOAD = json.dumps(
    {
        "titles": [{"title_raw": "Catan", "confidence": "high"}],
        "unidentified": [
            {
                "location": "middle shelf, between Catan and the frame edge",
                "partial_text": "EMP",
                "art_notes": "tall black box with gold lettering",
            }
        ],
    }
)


def test_unidentified_lands_in_artifact_and_summary(tmp_path, capsys):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf1.jpg")
    vision, _ = _vision_stub(OBJECT_PAYLOAD)
    run_extract(cfg, vision=vision)

    saved = json.loads(cfg.unidentified_path.read_text())
    assert saved == {
        "shelf1.jpg": [
            {
                "location": "middle shelf, between Catan and the frame edge",
                "partial_text": "EMP",
                "art_notes": "tall black box with gold lettering",
            }
        ]
    }
    out = capsys.readouterr().out
    assert "couldn't identify" in out
    assert "between Catan and the frame edge" in out


def test_clean_photo_writes_empty_unidentified(tmp_path):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf1.jpg")
    vision, _ = _vision_stub('{"titles": [{"title_raw": "Azul"}], "unidentified": []}')
    run_extract(cfg, vision=vision)
    assert json.loads(cfg.unidentified_path.read_text()) == {}


def test_bare_array_raw_cache_rebuilds(tmp_path):
    """A raw cache file may be a bare title array; it must still rebuild."""
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "old.jpg")
    raw_dir = cfg.data_dir / "extract_raw"
    raw_dir.mkdir(parents=True)
    (raw_dir / "old.jpg.json").write_text(json.dumps([_entry("Catan", "old.jpg")]))
    vision, calls = _vision_stub("[]")
    result = run_extract(cfg, vision=vision)
    assert calls == []  # cached photo untouched
    assert result[0]["title_raw"] == "Catan"
    assert json.loads(cfg.unidentified_path.read_text()) == {}


def test_force_reextracts_cached_photos(tmp_path):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf1.jpg")
    vision, calls = _vision_stub('[{"title_raw": "Azul"}]')
    run_extract(cfg, vision=vision)
    run_extract(cfg, vision=vision)
    assert len(calls) == 1  # cache hit

    vision2, calls2 = _vision_stub(OBJECT_PAYLOAD)
    result = run_extract(cfg, force=True, vision=vision2)
    assert len(calls2) == 1  # cache ignored
    assert result[0]["title_raw"] == "Catan"


def test_low_confidence_reads_are_surfaced(tmp_path, capsys):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf1.jpg")
    vision, _ = _vision_stub(
        '[{"title_raw": "Patchwork", "confidence": "medium"},'
        ' {"title_raw": "Catan", "confidence": "high"}]'
    )
    run_extract(cfg, vision=vision)
    out = capsys.readouterr().out
    assert "Low-confidence reads" in out
    assert "'Patchwork' (medium)" in out
    assert "'Catan'" not in out.split("Low-confidence reads")[1]


def test_corrupt_raw_cache_is_reextracted_not_skipped(tmp_path):
    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    _write_image(cfg.photos_dir / "shelf.jpg")
    vision, calls = _vision_stub('[{"title_raw": "Catan", "confidence": "high"}]')
    run_extract(cfg, vision=vision)
    raw = cfg.extract_raw_dir / "shelf.jpg.json"
    raw.write_text("{torn")  # corrupt the cache
    run_extract(cfg, vision=vision)
    assert len(calls) == 2  # re-extracted, not skipped
    json.loads(raw.read_text())  # cache healed


def test_systemic_failures_abort_and_exit_nonzero(tmp_path):
    import pytest
    import typer as _typer

    calls = []

    def broken_vision(image_b64, media_type):
        calls.append(1)
        raise RuntimeError("invalid x-api-key")

    cfg = _cfg(tmp_path)
    cfg.photos_dir.mkdir()
    for i in range(6):
        _write_image(cfg.photos_dir / f"p{i}.jpg")
    with pytest.raises(_typer.Exit):
        run_extract(cfg, vision=broken_vision)
    assert len(calls) == 3  # aborted after 3 identical failures