Model Picker — Rust source

//! Model Picker — filter + rank the AI model catalog.
//!
//! Language: Rust (edition 2021, standard library only)
//! Source:   CosmoDev polyglot showcase port of the Model Picker tool, ported
//!           from src/lib/modelPicker.ts (the canonical TypeScript
//!           implementation).
//! Live:     https://dev.cosmolabs.org/tools/model-picker
//! License:  display source — part of CosmoDev's polyglot tool pages.
//!
//! Design goals:
//!   - Pure + deterministic; never panics (no unwraps on data-derived floats —
//!     a NaN compares as Equal rather than poisoning the sort).
//!   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//!   - Self-contained: std only (no crates.io dependencies).
//!
//! Port notes: the TS lib imports `listModels` from src/lib/ai/models.ts and
//! lets `models` default to the bundled pricing snapshot
//! (src/data/ai-models.json). A dependency-free port cannot load that file, so
//! the filter step is inlined here and `models` is an explicit slice
//! parameter. Rust has no default arguments; the TS default `sort = 'price'`
//! is expressed by `SortKey::default()`. `Vec::sort_by` is stable, matching
//! `Array.prototype.sort`'s stability guarantee in the TS original — rows
//! comparing equal keep their input order.
//!
//! Every sort places null/missing values last and breaks ties on id ascending,
//! so output order is fully deterministic for a given model list.

use std::cmp::Ordering;

/// A model's qualitative class. Mirrors `Tier` in the TS lib.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Tier {
    Flagship,
    Balanced,
    Fast,
    Budget,
}

impl Tier {
    /// The TS wire value ("flagship" | "balanced" | "fast" | "budget").
    pub fn as_str(self) -> &'static str {
        match self {
            Tier::Flagship => "flagship",
            Tier::Balanced => "balanced",
            Tier::Fast => "fast",
            Tier::Budget => "budget",
        }
    }
}

/// Ranking selector. Mirrors the TS `SortKey` union; the TS default
/// `'price'` is carried by `#[default]`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SortKey {
    #[default]
    Price,
    Context,
    TokensPerDollar,
    Released,
}

/// One catalog entry. Mirrors the TS `AiModel` interface; `Option` plays the
/// role of TS `null` (unpriced / unknown).
#[derive(Debug, Clone, PartialEq)]
pub struct AiModel {
    pub id: String,
    pub name: String,
    pub vendor: String,
    pub family: String,
    pub tier: Tier,
    pub context_window: u64,
    pub max_output: u64,
    /// USD per 1M input tokens; `None` for open/unpriced models.
    pub input_per_m: Option<f64>,
    /// USD per 1M output tokens; `None` for open/unpriced models.
    pub output_per_m: Option<f64>,
    pub cache_read_per_m: Option<f64>,
    pub cache_write_per_m: Option<f64>,
    pub modalities: Vec<String>,
    pub knowledge_cutoff: Option<String>,
    /// ISO date string; `None` when unknown.
    pub released: Option<String>,
    pub open_weights: bool,
    pub reasoning: bool,
    pub tool_call: bool,
}

/// Mirrors the TS `ModelFilter` interface: every field is optional, and `None`
/// means "do not apply this constraint" (TS `undefined`).
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ModelFilter {
    pub vendor: Option<String>,
    pub tier: Option<Tier>,
    /// Minimum usable context window in tokens.
    pub min_context: Option<u64>,
    /// Maximum input price (USD per 1M tokens); unpriced models are skipped.
    pub max_input_per_m: Option<f64>,
    /// Required input modality, e.g. `"image"`.
    pub modality: Option<String>,
    /// Case-insensitive substring match on id, name, and vendor.
    pub search: Option<String>,
}

/// A curated entry point: a filter + the sort that makes that filter useful.
#[derive(Debug, Clone)]
pub struct Preset {
    pub filter: ModelFilter,
    pub sort: SortKey,
}

/// The TS `PRESETS` record's keys.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PresetKey {
    LongContext,
    CheapBulk,
    Flagship,
}

/// Curated entry points, mirroring the TS `PRESETS` record. A function (not a
/// `const`) because the values are built at runtime.
pub fn presets(key: PresetKey) -> Preset {
    match key {
        PresetKey::LongContext => Preset {
            filter: ModelFilter {
                min_context: Some(500_000),
                ..ModelFilter::default()
            },
            sort: SortKey::Context,
        },
        PresetKey::CheapBulk => Preset {
            filter: ModelFilter {
                max_input_per_m: Some(1.0),
                ..ModelFilter::default()
            },
            sort: SortKey::Price,
        },
        PresetKey::Flagship => Preset {
            filter: ModelFilter {
                tier: Some(Tier::Flagship),
                ..ModelFilter::default()
            },
            sort: SortKey::TokensPerDollar,
        },
    }
}

/// Output tokens per USD: 1e6 / output_per_m. `None` when unpriced.
/// Mirrors `tokensPerDollar()` in the TS lib.
pub fn tokens_per_dollar(m: &AiModel) -> Option<f64> {
    m.output_per_m.map(|p| 1_000_000.0 / p)
}

/// Filter the model list. Inline port of `listModels()` from
/// src/lib/ai/models.ts: `None` prices never satisfy `max_input_per_m`, and
/// search lowercases both sides before the substring check.
pub fn list_models(filter: &ModelFilter, models: &[AiModel]) -> Vec<AiModel> {
    let q = filter.search.as_ref().map(|s| s.to_lowercase());
    models
        .iter()
        .filter(|m| {
            if let Some(v) = &filter.vendor {
                if &m.vendor != v {
                    return false;
                }
            }
            if let Some(t) = filter.tier {
                if m.tier != t {
                    return false;
                }
            }
            if let Some(min) = filter.min_context {
                if m.context_window < min {
                    return false;
                }
            }
            if let Some(max) = filter.max_input_per_m {
                match m.input_per_m {
                    None => return false,
                    Some(p) if p > max => return false,
                    _ => {}
                }
            }
            if let Some(md) = &filter.modality {
                if !m.modalities.iter().any(|x| x == md) {
                    return false;
                }
            }
            if let Some(q) = &q {
                let hay = format!("{} {} {}", m.id, m.name, m.vendor).to_lowercase();
                if !hay.contains(q.as_str()) {
                    return false;
                }
            }
            true
        })
        .cloned()
        .collect()
}

/// id ascending — the shared stable tie-break for every sort.
fn by_id(a: &AiModel, b: &AiModel) -> Ordering {
    a.id.cmp(&b.id)
}

/// Float comparison that never panics: a NaN compares as Equal instead of
/// unwrapping a `partial_cmp`. Catalog prices are finite, so this only
/// matters for defensive robustness.
fn cmp_f64(a: f64, b: f64) -> Ordering {
    a.partial_cmp(&b).unwrap_or(Ordering::Equal)
}

/// price: input price ascending, unpriced last, id asc on ties.
fn compare_price(a: &AiModel, b: &AiModel) -> Ordering {
    match (&a.input_per_m, &b.input_per_m) {
        (None, None) => by_id(a, b),
        (None, Some(_)) => Ordering::Greater,
        (Some(_), None) => Ordering::Less,
        (Some(x), Some(y)) => cmp_f64(*x, *y).then_with(|| by_id(a, b)),
    }
}

/// context: context window descending, id asc on ties.
fn compare_context(a: &AiModel, b: &AiModel) -> Ordering {
    b.context_window
        .cmp(&a.context_window)
        .then_with(|| by_id(a, b))
}

/// tokensPerDollar: 1e6/output price descending, unpriced last, id asc on ties.
fn compare_tokens_per_dollar(a: &AiModel, b: &AiModel) -> Ordering {
    match (tokens_per_dollar(a), tokens_per_dollar(b)) {
        (None, None) => by_id(a, b),
        (None, Some(_)) => Ordering::Greater,
        (Some(_), None) => Ordering::Less,
        (Some(x), Some(y)) => cmp_f64(y, x).then_with(|| by_id(a, b)),
    }
}

/// released: newest first (ISO dates compare correctly as strings), unknown
/// last, id asc on ties.
fn compare_released(a: &AiModel, b: &AiModel) -> Ordering {
    match (&a.released, &b.released) {
        (None, None) => by_id(a, b),
        (None, Some(_)) => Ordering::Greater,
        (Some(_), None) => Ordering::Less,
        (Some(x), Some(y)) => y.cmp(x).then_with(|| by_id(a, b)),
    }
}

/// Sort comparator per key: value order first, `None` last, id asc on ties.
fn comparator(sort: SortKey) -> fn(&AiModel, &AiModel) -> Ordering {
    match sort {
        SortKey::Price => compare_price,
        SortKey::Context => compare_context,
        SortKey::TokensPerDollar => compare_tokens_per_dollar,
        SortKey::Released => compare_released,
    }
}

/// Filter the model list, then sort it. See `comparator` for the exact order.
pub fn pick_models(filter: &ModelFilter, sort: SortKey, models: &[AiModel]) -> Vec<AiModel> {
    let mut ranked = list_models(filter, models);
    ranked.sort_by(comparator(sort));
    ranked
}