Lumotia/crates/core/src/recommendation.rs

use crate::hardware::SystemProfile;
use crate::model_registry::{all_models, AccuracyTier, Engine, ModelEntry, SpeedTier};
use crate::types::Megabytes;

/// A model's suitability score for a given system. Higher is better.
/// No boolean flags — position in the ranked list conveys recommendation.
pub struct ScoredModel {
    pub entry: &'static ModelEntry,
    pub score: f64,
    pub reason: String,
}

/// Scores a single model against a system profile.
/// Pure function, no side effects.
pub fn score_model(model: &'static ModelEntry, profile: &SystemProfile) -> Option<ScoredModel> {
    if model.ram_required > profile.ram {
        return None;
    }

    let mut score = 0.0;
    let mut reasons: Vec<String> = Vec::new();

    score += match model.speed_tier {
        SpeedTier::Instant => 40.0,
        SpeedTier::Fast => 30.0,
        SpeedTier::Moderate => 20.0,
        SpeedTier::Slow => 10.0,
    };

    score += match model.accuracy_tier {
        AccuracyTier::Excellent => 30.0,
        AccuracyTier::Great => 20.0,
        AccuracyTier::Good => 10.0,
    };

    if let Some(gpu) = &profile.gpu {
        let has_accel = match model.engine {
            Engine::Whisper => {
                gpu.acceleration.metal || gpu.acceleration.vulkan || gpu.acceleration.cuda
            }
            Engine::Parakeet | Engine::Moonshine => {
                gpu.acceleration.cuda || gpu.acceleration.vulkan
            }
        };
        if has_accel {
            score += 15.0;
            reasons.push("GPU accelerated on your system".into());
        }
    }

    let headroom = Megabytes(profile.ram.0.saturating_sub(model.ram_required.0));
    if headroom > Megabytes::from_gb(4) {
        score += 10.0;
    }

    let reason = if reasons.is_empty() {
        model.description.to_string()
    } else {
        reasons.join(". ")
    };

    Some(ScoredModel {
        entry: model,
        score,
        reason,
    })
}

/// Scores all models and returns them ranked.
/// Index 0 is the recommendation. No flag arguments.
pub fn rank_recommendations(profile: &SystemProfile) -> Vec<ScoredModel> {
    let mut scored: Vec<ScoredModel> = all_models()
        .iter()
        .filter_map(|model| score_model(model, profile))
        .collect();

    scored.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    scored
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hardware::{CpuFeatures, CpuInfo, GpuAcceleration, GpuInfo, GpuVendor, Os};

    fn profile_with_ram(ram: Megabytes) -> SystemProfile {
        SystemProfile {
            ram,
            cpu: CpuInfo {
                logical_processors: 8,
                brand: "Test CPU".into(),
                features: CpuFeatures::default(),
            },
            gpu: None,
            os: Os::Windows,
        }
    }

    fn profile_with_gpu(ram: Megabytes) -> SystemProfile {
        SystemProfile {
            ram,
            cpu: CpuInfo {
                logical_processors: 8,
                brand: "Test CPU".into(),
                features: CpuFeatures::default(),
            },
            gpu: Some(GpuInfo {
                vendor: GpuVendor::Nvidia,
                vram: Megabytes(8192),
                acceleration: GpuAcceleration {
                    cuda: true,
                    metal: false,
                    vulkan: true,
                },
            }),
            os: Os::Windows,
        }
    }

    #[test]
    fn score_model_excludes_models_exceeding_available_ram() {
        let profile = profile_with_ram(Megabytes(256));
        let model = all_models()
            .iter()
            .find(|m| m.ram_required > Megabytes(256))
            .expect("need a model larger than 256 MB");

        let result = score_model(model, &profile);

        assert!(result.is_none());
    }

    #[test]
    fn score_model_includes_models_fitting_in_ram() {
        let profile = profile_with_ram(Megabytes(16384));
        let model = &all_models()[0];

        let result = score_model(model, &profile);

        assert!(result.is_some());
    }

    #[test]
    fn score_model_boosts_gpu_accelerated_models() {
        let model = all_models()
            .iter()
            .find(|m| m.engine == Engine::Parakeet)
            .expect("need a Parakeet model");

        let gpu_score = score_model(model, &profile_with_gpu(Megabytes(16384)))
            .unwrap()
            .score;
        let cpu_score = score_model(model, &profile_with_ram(Megabytes(16384)))
            .unwrap()
            .score;

        assert!(gpu_score > cpu_score);
    }

    #[test]
    fn rank_recommendations_places_highest_score_first() {
        let profile = profile_with_ram(Megabytes(16384));

        let ranked = rank_recommendations(&profile);

        assert!(ranked.len() >= 2);
        assert!(ranked[0].score >= ranked[1].score);
    }

    #[test]
    fn rank_recommendations_returns_empty_for_very_low_ram() {
        let profile = profile_with_ram(Megabytes(128));

        let ranked = rank_recommendations(&profile);

        assert!(ranked.is_empty());
    }

    #[test]
    fn parakeet_is_top_recommendation_when_hardware_supports_it() {
        // Any machine that fits Parakeet in RAM should see it ranked first —
        // Parakeet-TDT is English-only but beats Whisper on English at lower
        // latency, so it's Lumotia's default recommendation when eligible.
        // (Users on non-English languages adjust manually — handled at the
        // settings-UI level, not at the scoring level for now.)
        let profile = profile_with_ram(Megabytes(16384));

        let ranked = rank_recommendations(&profile);
        let top = ranked.first().expect("at least one model ranks");

        assert_eq!(top.entry.engine, Engine::Parakeet);
    }
}