Use policy scoring weights in release search (#74)

crates/arr-api/src/search.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 use std::time::UNIX_EPOCH;
 
 use arr_core::policy::{evaluate, Candidate};
+use arr_core::score::score;
 use arr_core::{
     DolbyVisionProfile, HdrRules, Language, Policy, PolicyId, RequiredAudio, Resolution, Rule,
     ScoreWeights, SizeBand, Source, TitleOverrides, Verdict,
@@ -80,6 +81,7 @@ struct PolicyRow {
     size_bands: String,
     resolution_pref: String,
     source_weights: String,
+    score_weights: String,
 }
 
 #[derive(Debug, Deserialize)]
@@ -102,6 +104,13 @@ struct SizeBandJson {
     penalty_points_per_gib_over: i32,
 }
 
+#[derive(Debug, Deserialize)]
+struct ScoreWeightsJson {
+    size_at_target: i32,
+    source_tier: i32,
+    seeder_doubling: i32,
+}
+
 #[derive(Debug, Deserialize)]
 struct OverridesJson {
     #[serde(default)]
@@ -254,7 +263,7 @@ pub async fn releases(
     Query(query): Query<ReleasesQuery>,
 ) -> Result<Json<Vec<ClassifiedRelease>>, ApiError> {
     let database = state.database().ok_or(ApiError::Unavailable)?;
-    let movie = sqlx::query!(r#"SELECT m.title AS "title!: String", m.tmdb_id AS "tmdb_id!: i64", m.original_language, m.overrides AS "overrides!: serde_json::Value", p.id AS "policy_id!: i64", p.name AS "policy_name!: String", p.required_audio AS "required_audio!: String", p.dub_blacklist AS "dub_blacklist!: String", p.hdr_rules AS "hdr_rules!: String", p.size_bands AS "size_bands!: String", p.resolution_pref AS "resolution_pref!: String", p.source_weights AS "source_weights!: String" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.id = ?"#, query.movie_id)
+    let movie = sqlx::query!(r#"SELECT m.title AS "title!: String", m.tmdb_id AS "tmdb_id!: i64", m.original_language, m.overrides AS "overrides!: serde_json::Value", p.id AS "policy_id!: i64", p.name AS "policy_name!: String", p.required_audio AS "required_audio!: String", p.dub_blacklist AS "dub_blacklist!: String", p.hdr_rules AS "hdr_rules!: String", p.size_bands AS "size_bands!: String", p.resolution_pref AS "resolution_pref!: String", p.source_weights AS "source_weights!: String", p.score_weights AS "score_weights!: String" FROM movies m JOIN roots r ON r.id = m.root_id JOIN policies p ON p.id = r.policy_id WHERE m.id = ?"#, query.movie_id)
         .fetch_optional(database.pool()).await?.ok_or(ApiError::NotFound)?;
 
     let tmdb = tmdb_client(&state)?
@@ -290,6 +299,7 @@ pub async fn releases(
         size_bands: movie.size_bands,
         resolution_pref: movie.resolution_pref,
         source_weights: movie.source_weights,
+        score_weights: movie.score_weights,
     })?;
     let overrides: OverridesJson = serde_json::from_value(movie.overrides)
         .map_err(|error| ApiError::Database(error.to_string()))?;
@@ -405,6 +415,7 @@ fn policy_from_row(row: PolicyRow) -> Result<Policy, ApiError> {
     let bands: BTreeMap<String, SizeBandJson> = json(&row.size_bands)?;
     let resolutions: Vec<String> = json(&row.resolution_pref)?;
     let weights: BTreeMap<String, i32> = json(&row.source_weights)?;
+    let score_weights: ScoreWeightsJson = json(&row.score_weights)?;
     Ok(Policy {
         id: PolicyId(row.policy_id),
         name: row.policy_name,
@@ -451,9 +462,11 @@ fn policy_from_row(row: PolicyRow) -> Result<Policy, ApiError> {
             .into_iter()
             .filter_map(|(source, weight)| source_value(&source).map(|source| (source, weight)))
             .collect(),
-        // The row's score_weights column is not selected here yet; the seeded
-        // values match these defaults. Wiring the column through is #68.
-        score_weights: ScoreWeights::default(),
+        score_weights: ScoreWeights {
+            size_at_target: score_weights.size_at_target,
+            source_tier: score_weights.source_tier,
+            seeder_doubling: score_weights.seeder_doubling,
+        },
     })
 }
 
@@ -472,7 +485,17 @@ fn classify(
         release.size,
     );
     let (verdict, rule) = verdict(&evaluation.verdict);
-    let score = score(policy, &parsed, release.size, release.seeders);
+    let score = score(
+        policy,
+        Candidate::PreGrab(&parsed),
+        release.size.unwrap_or_default(),
+        release.seeders.unwrap_or_default(),
+    );
+    let score = if release.size.is_some() {
+        score.total
+    } else {
+        score.source.saturating_add(score.seeders)
+    };
     Ok(ClassifiedRelease {
         indexer_id: release.indexer_id,
         guid: release.guid,
@@ -495,44 +518,6 @@ fn classify(
     })
 }
 
-fn score(
-    policy: &Policy,
-    parsed: &arr_parse::NameClaims,
-    size: Option<u64>,
-    seeders: Option<u32>,
-) -> i64 {
-    let resolution = parsed.resolution.map(Resolution::from);
-    let resolution_score = resolution
-        .and_then(|value| {
-            policy
-                .resolution_preference
-                .iter()
-                .position(|candidate| *candidate == value)
-        })
-        .map_or(0, |position| {
-            10_000 - i64::try_from(position).unwrap_or(0) * 5_000
-        });
-    let size_score = resolution
-        .and_then(|value| policy.size_bands.get(&value))
-        .zip(size)
-        .map_or(0, |(band, bytes)| {
-            if bytes < band.floor_bytes {
-                -i64::try_from((band.floor_bytes - bytes) / 100_000_000).unwrap_or(i64::MAX)
-            } else {
-                -i64::try_from(bytes.abs_diff(band.target_bytes) / 100_000_000).unwrap_or(i64::MAX)
-            }
-        });
-    let source_score = i64::from(
-        parsed
-            .source
-            .map(Source::from)
-            .and_then(|source| policy.source_weights.get(&source).copied())
-            .unwrap_or(0),
-    ) * 5;
-    let seeder_score = seeders.map_or(0, |count| i64::from(count.saturating_add(1).ilog2()) * 2);
-    resolution_score + size_score + source_score + seeder_score
-}
-
 fn verdict(verdict: &Verdict) -> (&'static str, Option<String>) {
     match verdict {
         Verdict::Eligible => ("eligible", None),
@@ -729,6 +714,12 @@ mod tests {
         let (_dir, state, base) = application(&tmdb, &prowlarr).await;
         sqlx::query("INSERT INTO movies (tmdb_id, title, year, original_language, root_id) VALUES (693134, 'Dune Part Two', 2024, 'en', 2)")
             .execute(state.database().expect("database").pool()).await.expect("movie");
+        sqlx::query(
+            "UPDATE policies SET score_weights = '{\"size_at_target\":0,\"source_tier\":0,\"seeder_doubling\":0}'",
+        )
+        .execute(state.database().expect("database").pool())
+        .await
+        .expect("score weights");
 
         let response = reqwest::get(format!("{base}/api/releases?movie_id=1"))
             .await
@@ -744,7 +735,87 @@ mod tests {
             .find(|release| release["verdict"] == "rejected")
             .expect("rejected");
         assert_eq!(rejected["rule"], "source");
-        assert!(releases[0]["score"].as_i64().is_some());
+        let eligible = releases
+            .iter()
+            .find(|release| release["guid"] == "good")
+            .expect("eligible");
+        assert_eq!(eligible["score"], 0);
+    }
+
+    #[test]
+    fn policy_row_uses_persisted_score_weights() {
+        let policy = policy_from_row(PolicyRow {
+            policy_id: 1,
+            policy_name: "test".into(),
+            required_audio: r#"{"require":"original_language"}"#.into(),
+            dub_blacklist: "[]".into(),
+            hdr_rules: "{}".into(),
+            size_bands:
+                r#"{"2160p":{"floor_gib":8,"target_gib":22,"penalty_points_per_gib_over":60}}"#
+                    .into(),
+            resolution_pref: r#"["2160p"]"#.into(),
+            source_weights: r#"{"WEB-DL":2}"#.into(),
+            score_weights: r#"{"size_at_target":2000,"source_tier":7,"seeder_doubling":11}"#.into(),
+        })
+        .expect("policy row");
+
+        assert_eq!(
+            policy.score_weights,
+            ScoreWeights {
+                size_at_target: 2000,
+                source_tier: 7,
+                seeder_doubling: 11,
+            }
+        );
+    }
+
+    #[test]
+    fn releases_without_sizes_skip_the_size_score() {
+        let policy = Policy {
+            id: PolicyId(1),
+            name: "test".into(),
+            required_audio: RequiredAudio::OriginalLanguage,
+            dub_blacklist: Vec::new(),
+            hdr_rules: HdrRules {
+                rejected_dolby_vision_profiles: Vec::new(),
+            },
+            size_bands: std::collections::BTreeMap::from([(
+                Resolution::R2160p,
+                SizeBand {
+                    floor_bytes: gib(8),
+                    target_bytes: gib(22),
+                    penalty_points_per_gib_over: 60,
+                },
+            )]),
+            resolution_preference: vec![Resolution::R2160p],
+            source_weights: std::collections::BTreeMap::from([(Source::WebDl, 2)]),
+            score_weights: ScoreWeights::default(),
+        };
+        let release = SearchRelease {
+            indexer_id: 1,
+            guid: "release".into(),
+            name: "Dune.Part.Two.2024.2160p.WEB-DL".into(),
+            size: None,
+            seeders: Some(8),
+            publish_date: None,
+            download_url: "https://tracker/release".into(),
+        };
+        let parsed = arr_parse::parse(&release.name);
+        let core_score = score(&policy, Candidate::PreGrab(&parsed), 0, 8);
+
+        let classified = classify(
+            release,
+            &policy,
+            &MovieOverrides::default(),
+            &Language::Other("en".into()),
+        )
+        .expect("classified release");
+
+        assert_eq!(
+            classified.score,
+            core_score.source.saturating_add(core_score.seeders)
+        );
+        assert_ne!(classified.score, core_score.total);
     }
 
     #[test]