malvin 0.2.7

Non-interactive research and coding agent
use std::collections::{HashMap, HashSet};

use pi::sdk::ModelRegistry;

use super::models_list::PiModelListing;

pub(crate) fn merge_registry_with_live(
    registry: &ModelRegistry,
    live_by_provider: &HashMap<String, Vec<String>>,
) -> Vec<PiModelListing> {
    let static_by_key = static_registry_lookup(registry);
    let mut out = Vec::new();
    let mut seen = HashSet::new();
    append_live_models(&mut out, &mut seen, &static_by_key, live_by_provider);
    append_static_models_without_live(&mut out, &mut seen, registry, live_by_provider);
    out
}

fn static_registry_lookup(
    registry: &ModelRegistry,
) -> HashMap<(String, String), &pi::models::ModelEntry> {
    registry
        .models()
        .iter()
        .map(|entry| {
            (
                (entry.model.provider.clone(), entry.model.id.clone()),
                entry,
            )
        })
        .collect()
}

fn append_live_models(
    out: &mut Vec<PiModelListing>,
    seen: &mut HashSet<String>,
    static_by_key: &HashMap<(String, String), &pi::models::ModelEntry>,
    live_by_provider: &HashMap<String, Vec<String>>,
) {
    for (provider, ids) in live_by_provider {
        for id in ids {
            let full_id = format!("{provider}/{id}");
            if !seen.insert(full_id.clone()) {
                continue;
            }
            let (name, thinking) = static_by_key
                .get(&(provider.clone(), id.clone()))
                .map_or_else(
                    || (id.clone(), None),
                    |entry| (entry.model.name.clone(), Some(entry.model.reasoning)),
                );
            out.push(PiModelListing {
                id: full_id,
                name,
                thinking,
            });
        }
    }
}

fn append_static_models_without_live(
    out: &mut Vec<PiModelListing>,
    seen: &mut HashSet<String>,
    registry: &ModelRegistry,
    live_by_provider: &HashMap<String, Vec<String>>,
) {
    for entry in registry.models() {
        let provider = entry.model.provider.as_str();
        if live_by_provider.contains_key(provider) {
            continue;
        }
        let full_id = format!("{provider}/{}", entry.model.id);
        if seen.insert(full_id.clone()) {
            out.push(PiModelListing {
                id: full_id,
                name: entry.model.name.clone(),
                thinking: Some(entry.model.reasoning),
            });
        }
    }
}