From d28fb0a0a7fa54566c4ad2d72cc1c825adb462ae Mon Sep 17 00:00:00 2001
From: Ben Stull <ben.stull@wiggleverse.org>
Date: Thu, 4 Jun 2026 01:35:38 -0700
Subject: [PATCH] feat: selection mode filtering with av fallback

---
 hef/selection.py        | 18 ++++++++++++++++++
 tests/test_selection.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/hef/selection.py b/hef/selection.py
index 4a0c6ff..3a8d5f3 100644
--- a/hef/selection.py
+++ b/hef/selection.py
@@ -35,3 +35,21 @@ def distance(a: Coordinate, b: Coordinate, weights: Weights = Weights()) -> floa
 def record_coordinate(record: Record) -> Coordinate:
     """The (left, right, dark, light) coordinate of a catalog record."""
     return Coordinate(record.left, record.right, record.dark, record.light)
+
+
+def candidates_for_mode(records, mode: str, pool_size: int) -> list[Record]:
+    """Records eligible for a content mode.
+
+    For 'av', if fewer than pool_size native 'av' records exist, fall back to
+    including 'audio' and 'video' records so the pool is never starved.
+    """
+    if mode not in CONTENT_MODES:
+        raise ValueError(
+            f"candidates_for_mode expects a content mode {sorted(CONTENT_MODES)}, "
+            f"got {mode!r}"
+        )
+    primary = [r for r in records if r.mode == mode]
+    if mode == "av" and len(primary) < pool_size:
+        extra = [r for r in records if r.mode in {"audio", "video"}]
+        return primary + extra
+    return primary
diff --git a/tests/test_selection.py b/tests/test_selection.py
index ac6d1d4..8012986 100644
--- a/tests/test_selection.py
+++ b/tests/test_selection.py
@@ -48,3 +48,34 @@ def test_weights_scale_planes_independently():
 def test_record_coordinate_extracts_axes():
     record = make_record(left=1, right=2, dark=3, light=4)
     assert record_coordinate(record) == Coordinate(1, 2, 3, 4)
+
+
+from hef.selection import candidates_for_mode
+
+
+def test_candidates_filter_to_exact_mode():
+    records = [make_record(id="v", mode="video"), make_record(id="a", mode="audio")]
+    result = candidates_for_mode(records, "audio", pool_size=4)
+    assert [r.id for r in result] == ["a"]
+
+
+def test_av_does_not_fall_back_when_enough_av():
+    records = [make_record(id=f"av{i}", mode="av") for i in range(4)]
+    records.append(make_record(id="a", mode="audio"))
+    result = candidates_for_mode(records, "av", pool_size=4)
+    assert {r.id for r in result} == {"av0", "av1", "av2", "av3"}
+
+
+def test_av_falls_back_to_audio_and_video_when_thin():
+    records = [
+        make_record(id="av0", mode="av"),
+        make_record(id="a", mode="audio"),
+        make_record(id="v", mode="video"),
+    ]
+    result = candidates_for_mode(records, "av", pool_size=4)
+    assert {r.id for r in result} == {"av0", "a", "v"}
+
+
+def test_candidates_rejects_none_mode():
+    with pytest.raises(ValueError):
+        candidates_for_mode([], "none", pool_size=4)