//! models.dev metadata enrichment (multi-provider expansion).
//!
//! Most provider `/models` listings return bare ids — no context window, no
//! thinking levels. This module carries per-model metadata keyed by
//! models.dev provider id, sourced from the bundled snapshot
//! (`enrichment_snapshot.json`, regenerated by
//! `scripts/gen_enrichment_snapshot.py`) or a runtime refresh fetched by the
//! shell. It NEVER invents model availability: the live listing is the only
//! source of which models exist — enrichment only fills metadata gaps, and
//! wire-served values always win (see [`enrich_wire_model`]).

use std::collections::BTreeMap;
use std::sync::LazyLock;

/// One model's enrichment metadata. All fields optional-by-default so the
/// snapshot stays minimal.
#[derive(Debug, Clone, Default, PartialEq, serde::Deserialize, serde::Serialize)]
pub struct EnrichmentModel {
    /// Max context window in tokens (`limit.context`).
    #[serde(default, skip_serializing_if = "is_zero")]
    pub context: u64,
    /// Max output tokens (`limit.output`); fills a wire-unserved output cap
    /// (Anthropic 400s when `max_tokens` exceeds the model's limit).
    #[serde(default, skip_serializing_if = "is_zero")]
    pub output: u64,
    /// Model supports reasoning/thinking.
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub reasoning: bool,
    /// Selectable effort levels (canonical tokens, e.g. ["low","high","max"]).
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub efforts: Vec<String>,
    /// Accepts image input.
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub image_in: bool,
    /// Supports tool calling (used later to filter non-agentic listings).
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub tool_call: bool,
    /// Human display name.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

fn is_zero(v: &u64) -> bool {
    *v == 0
}

/// `models.dev provider id -> model id -> metadata`.
pub type EnrichmentCatalog = BTreeMap<String, BTreeMap<String, EnrichmentModel>>;

/// The bundled snapshot, embedded at compile time — RAW models.dev shape,
/// filtered to the providers kigi references (pure-filter script:
/// `scripts/gen_enrichment_snapshot.py`). OFFLINE FALLBACK for the runtime
/// refresh; both parse through [`parse_api_json`] so there is exactly one
/// field interpretation.
pub const ENRICHMENT_SNAPSHOT_JSON: &str = include_str!("../enrichment_snapshot.json");

static BUNDLED: LazyLock<EnrichmentCatalog> = LazyLock::new(|| {
    // Baked-in JSON — a mismatch here is a developer error, not a runtime
    // condition (same policy as default_models.json).
    parse_api_json(ENRICHMENT_SNAPSHOT_JSON, None)
        .expect("enrichment_snapshot.json: invalid JSON (regenerate via script)")
});

// ── Raw models.dev api.json shape (parse-only) ──────────────────────────────

#[derive(serde::Deserialize)]
struct RawProvider {
    #[serde(default)]
    models: BTreeMap<String, RawModel>,
}

#[derive(serde::Deserialize, Default)]
#[serde(default)]
struct RawModel {
    name: Option<String>,
    reasoning: bool,
    reasoning_options: Vec<RawReasoningOption>,
    limit: RawLimit,
    modalities: RawModalities,
    tool_call: bool,
}

#[derive(serde::Deserialize, Default)]
#[serde(default)]
struct RawReasoningOption {
    r#type: String,
    values: Vec<String>,
}

#[derive(serde::Deserialize, Default)]
#[serde(default)]
struct RawLimit {
    context: u64,
    output: u64,
}

#[derive(serde::Deserialize, Default)]
#[serde(default)]
struct RawModalities {
    input: Vec<String>,
}

/// Parse a models.dev `api.json` document (full download or the bundled
/// filtered snapshot) into the in-memory catalog. `keep`: restrict to these
/// provider ids (runtime refresh filters the full 3MB download to the
/// registry's providers); `None` keeps everything present.
///
/// Blast-radius confinement: providers are filtered on RAW keys first and
/// typed-parsed individually — a schema drift in one of the ~70 providers
/// kigi never keeps cannot fail the whole refresh, and a malformed KEPT
/// provider is warn-skipped (its models fall to defaults) rather than
/// killing the others. Only a document that isn't a JSON object errors
/// (the caller falls back to cache/bundled — never a silently empty
/// catalog). Keys starting with `_` (provenance stamps) are skipped.
pub fn parse_api_json(
    json: &str,
    keep: Option<&std::collections::BTreeSet<&str>>,
) -> Result<EnrichmentCatalog, serde_json::Error> {
    let raw: BTreeMap<String, serde_json::Value> = serde_json::from_str(json)?;
    let mut catalog = EnrichmentCatalog::new();
    for (pid, value) in raw {
        if pid.starts_with('_') {
            continue;
        }
        if let Some(keep) = keep
            && !keep.contains(pid.as_str())
        {
            continue;
        }
        let provider: RawProvider = match serde_json::from_value(value) {
            Ok(p) => p,
            Err(e) => {
                tracing::warn!(provider = %pid, error = %e,
                    "models.dev provider entry malformed; skipping (models fall to defaults)");
                continue;
            }
        };
        let models = provider
            .models
            .into_iter()
            .map(|(mid, m)| {
                let efforts = m
                    .reasoning_options
                    .into_iter()
                    .find(|o| o.r#type == "effort" && !o.values.is_empty())
                    .map(|o| o.values)
                    .unwrap_or_default();
                let meta = EnrichmentModel {
                    context: m.limit.context,
                    output: m.limit.output,
                    reasoning: m.reasoning,
                    efforts,
                    image_in: m.modalities.input.iter().any(|s| s == "image"),
                    tool_call: m.tool_call,
                    name: m.name,
                };
                (mid, meta)
            })
            .collect();
        catalog.insert(pid, models);
    }
    Ok(catalog)
}

/// The compiled-in enrichment catalog.
pub fn bundled_enrichment() -> &'static EnrichmentCatalog {
    &BUNDLED
}

/// Look up one model's metadata. `None` when the provider or model is
/// unknown to the catalog (callers fall through to defaults).
pub fn lookup<'a>(
    catalog: &'a EnrichmentCatalog,
    models_dev_id: &str,
    model_id: &str,
) -> Option<&'a EnrichmentModel> {
    catalog.get(models_dev_id)?.get(model_id)
}

/// Fill metadata gaps on a live listing entry. WIRE WINS: a field the
/// provider served is never overwritten — enrichment only supplies what the
/// wire left absent/zero. Never changes the model id (availability stays
/// wire-truth).
pub fn enrich_wire_model(wire: &mut crate::WireModel, meta: &EnrichmentModel) {
    if wire.context_length == 0 && meta.context > 0 {
        wire.context_length = meta.context;
    }
    if wire.max_output_tokens == 0 && meta.output > 0 {
        wire.max_output_tokens = meta.output;
    }
    if meta.reasoning {
        wire.supports_reasoning = true;
    }
    if meta.image_in {
        wire.supports_image_in = true;
    }
    if wire.display_name.is_none() {
        wire.display_name = meta.name.clone();
    }
    if wire.think_efforts.is_none() && !meta.efforts.is_empty() {
        wire.think_efforts = Some(crate::WireThinkEfforts {
            support: true,
            valid_efforts: meta.efforts.clone(),
            // The provider's implicit default applies when the user doesn't
            // pick a level; models.dev doesn't record one.
            default_effort: None,
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The bundled snapshot parses and covers the providers this expansion
    /// references; kimi entries cross-check against the live wire values the
    /// registry already knows (guards against a corrupted regeneration).
    #[test]
    fn bundled_snapshot_parses_and_cross_checks_kimi() {
        let catalog = bundled_enrichment();
        assert!(
            catalog.len() >= 25,
            "snapshot lost providers: {}",
            catalog.len()
        );
        let k3 = lookup(catalog, "kimi-for-coding", "k3").expect("k3 present");
        assert_eq!(k3.context, 1_048_576, "k3 context must match the live wire");
        assert_eq!(k3.efforts, ["low", "high", "max"]);
        assert!(k3.reasoning);
        let opus = lookup(catalog, "anthropic", "claude-opus-4-8").expect("opus present");
        assert_eq!(opus.context, 1_000_000);
        assert_eq!(opus.efforts, ["low", "medium", "high", "xhigh", "max"]);
    }

    /// Wire-served fields are never overwritten; absent fields are filled.
    #[test]
    fn enrich_fills_gaps_and_never_overwrites_wire() {
        let meta = EnrichmentModel {
            context: 400_000,
            output: 64_000,
            reasoning: true,
            efforts: vec!["low".into(), "high".into()],
            image_in: true,
            name: Some("GPT Test".into()),
            ..Default::default()
        };
        // Bare listing entry (OpenAI-style: id only).
        let mut bare: crate::WireModel =
            serde_json::from_value(serde_json::json!({ "id": "gpt-test" })).unwrap();
        enrich_wire_model(&mut bare, &meta);
        assert_eq!(bare.context_length, 400_000);
        assert_eq!(bare.max_output_tokens, 64_000, "output cap filled");
        assert!(bare.supports_reasoning);
        assert!(bare.supports_image_in);
        assert_eq!(bare.display_name.as_deref(), Some("GPT Test"));
        let efforts = bare.think_efforts.expect("efforts filled");
        assert!(efforts.support);
        assert_eq!(efforts.valid_efforts, ["low", "high"]);
        assert_eq!(efforts.default_effort, None);

        // Wire-served entry: nothing may change.
        let mut served: crate::WireModel = serde_json::from_value(serde_json::json!({
            "id": "gpt-test",
            "context_length": 123,
            "max_output_tokens": 77,
            "display_name": "Wire Name",
            "think_efforts": { "support": true, "valid_efforts": ["max"] }
        }))
        .unwrap();
        enrich_wire_model(&mut served, &meta);
        assert_eq!(served.context_length, 123, "wire context wins");
        assert_eq!(served.max_output_tokens, 77, "wire output cap wins");
        assert_eq!(served.display_name.as_deref(), Some("Wire Name"));
        assert_eq!(
            served.think_efforts.unwrap().valid_efforts,
            ["max"],
            "wire efforts win"
        );
    }

    /// The runtime refresh parses the FULL api.json (extra fields like cost/
    /// env/doc present) and filters to the registry's provider ids.
    #[test]
    fn parse_api_json_filters_and_tolerates_unknown_fields() {
        let full = serde_json::json!({
            "openai": {
                "id": "openai", "env": ["OPENAI_API_KEY"], "doc": "https://x",
                "models": {
                    "gpt-test": {
                        "name": "GPT Test",
                        "reasoning": true,
                        "reasoning_options": [
                            {"type": "effort", "values": ["low", "high", "max"]}
                        ],
                        "limit": {"context": 400000, "output": 128000},
                        "modalities": {"input": ["text", "image"], "output": ["text"]},
                        "tool_call": true,
                        "cost": {"input": 1.25, "output": 10}
                    }
                }
            },
            "unwanted-provider": { "models": { "m": {} } }
        })
        .to_string();
        let keep: std::collections::BTreeSet<&str> = ["openai"].into();
        let catalog = parse_api_json(&full, Some(&keep)).unwrap();
        assert!(!catalog.contains_key("unwanted-provider"));
        let m = lookup(&catalog, "openai", "gpt-test").unwrap();
        assert_eq!(m.context, 400_000);
        assert_eq!(m.output, 128_000);
        assert_eq!(m.efforts, ["low", "high", "max"]);
        assert!(m.reasoning && m.image_in && m.tool_call);
        assert_eq!(m.name.as_deref(), Some("GPT Test"));
        // Malformed document errors — callers must fall back loudly, never
        // proceed with a silently empty catalog.
        assert!(parse_api_json("not json", None).is_err());
    }

    /// Every registry `models_dev_id` must be covered by the bundled
    /// snapshot — a registry row added without updating the script's
    /// TARGETS would silently diverge bundled-vs-refresh behavior.
    #[test]
    fn bundled_snapshot_covers_every_registry_models_dev_id() {
        let catalog = bundled_enrichment();
        for platform in crate::PlatformId::ALL {
            if let Some(dev_id) = platform.models_dev_id() {
                assert!(
                    catalog.contains_key(dev_id),
                    "{}: models_dev_id {dev_id:?} missing from the bundled \
                     snapshot — add it to scripts/gen_enrichment_snapshot.py \
                     TARGETS and regenerate",
                    platform.as_str(),
                );
            }
        }
    }

    /// Field-coverage guard against script/parser drift: for every field
    /// `RawModel` reads, at least one bundled model must carry a non-default
    /// value. A regeneration that dropped a MODEL_KEYS entry (e.g.
    /// `modalities`) would zero that field across the whole snapshot and
    /// fail here, instead of silently diverging from runtime refreshes.
    #[test]
    fn bundled_snapshot_carries_every_parsed_field() {
        let all: Vec<&EnrichmentModel> = bundled_enrichment()
            .values()
            .flat_map(|models| models.values())
            .collect();
        assert!(all.iter().any(|m| m.context > 0), "no context anywhere");
        assert!(all.iter().any(|m| m.output > 0), "no output anywhere");
        assert!(all.iter().any(|m| m.reasoning), "no reasoning anywhere");
        assert!(
            all.iter().any(|m| !m.efforts.is_empty()),
            "no efforts anywhere"
        );
        assert!(all.iter().any(|m| m.image_in), "no image_in anywhere");
        assert!(all.iter().any(|m| m.tool_call), "no tool_call anywhere");
        assert!(all.iter().any(|m| m.name.is_some()), "no names anywhere");
    }

    /// One malformed provider (schema drift on models.dev) must not kill
    /// the whole refresh — kept siblings still parse; only a non-object
    /// document errors.
    #[test]
    fn malformed_provider_is_skipped_not_fatal() {
        let doc = serde_json::json!({
            "good": { "models": { "m": { "limit": {"context": 7} } } },
            "drifted": { "models": "this is not an object" },
            "_meta": { "source": "stamp, must be ignored" }
        })
        .to_string();
        let catalog = parse_api_json(&doc, None).expect("document parses");
        assert_eq!(
            lookup(&catalog, "good", "m").map(|m| m.context),
            Some(7),
            "sibling providers must survive one drifted provider"
        );
        assert!(!catalog.contains_key("drifted"));
        assert!(!catalog.contains_key("_meta"));
    }

    #[test]
    fn lookup_misses_are_none() {
        let catalog = bundled_enrichment();
        assert!(lookup(catalog, "no-such-provider", "x").is_none());
        assert!(lookup(catalog, "openai", "no-such-model").is_none());
    }
}
