olai_codegen/

openapi_enrich.rs

use std::collections::{BTreeMap, HashSet};
use std::fs;
use std::path::Path;

use glob::glob;
use protobuf::Message;
use serde_json::Map as JsonMap;
use serde_yaml::Value as YamlValue;

use crate::Result;
use crate::parsing::{CodeGenMetadata, parse_file_descriptor_set};
use crate::utils::last_segment;

/// Run Pass 0 (gnostic ref fix), Pass 1 (validation enrichment), optionally Pass 2
/// (path/body dedup), and Pass 3 (schema-key aliasing).
pub fn run(
    spec: &Path,
    jsonschema_dir: &Path,
    camel_case: bool,
    descriptors: Option<&Path>,
    schema_renames: &BTreeMap<String, String>,
    schema_overrides: &HashSet<String>,
) -> Result<()> {
    let spec_str = fs::read_to_string(spec).map_err(|e| file_error("read", spec, e))?;
    let mut doc: YamlValue =
        serde_yaml::from_str(&spec_str).map_err(|e| file_error("parse YAML", spec, e))?;

    fix_gnostic_refs(&mut doc);
    enrich_from_jsonschema(&mut doc, jsonschema_dir, camel_case, schema_overrides)?;

    if let Some(desc_path) = descriptors {
        let bytes =
            fs::read(desc_path).map_err(|e| file_error("read descriptors", desc_path, e))?;
        let fds = protobuf::descriptor::FileDescriptorSet::parse_from_bytes(&bytes)
            .map_err(|e| crate::Error::Build(format!("Failed to parse descriptors: {e}")))?;
        let metadata = parse_file_descriptor_set(&fds)?;
        dedup_path_params(&mut doc, &metadata);
    }

    rename_schemas(&mut doc, schema_renames);

    let out = serde_yaml::to_string(&doc)
        .map_err(|e| crate::Error::Build(format!("Failed to serialize YAML: {e}")))?;
    fs::write(spec, out).map_err(|e| file_error("write", spec, e))?;

    Ok(())
}

/// `Failed to <action> <path>: <error>`.
fn file_error(action: &str, path: &Path, error: impl std::fmt::Display) -> crate::Error {
    crate::Error::Build(format!("Failed to {action} {}: {error}", path.display()))
}

// ── Pass 0 ────────────────────────────────────────────────────────────────────

/// Rewrite gnostic-generated `#/$defs/a.b.v1.TypeName.schema.strict.json` refs to
/// valid OpenAPI `#/components/schemas/TypeName` refs throughout the document.
fn fix_gnostic_refs(doc: &mut YamlValue) {
    match doc {
        YamlValue::Mapping(map) => {
            if let Some(YamlValue::String(s)) = map.get_mut("$ref")
                && let Some(fixed) = rewrite_gnostic_ref(s)
            {
                *s = fixed;
            }
            for (_, v) in map.iter_mut() {
                fix_gnostic_refs(v);
            }
        }
        YamlValue::Sequence(seq) => {
            for item in seq.iter_mut() {
                fix_gnostic_refs(item);
            }
        }
        _ => {}
    }
}

fn rewrite_gnostic_ref(ref_str: &str) -> Option<String> {
    let stem = ref_str
        .strip_prefix("#/$defs/")?
        .strip_suffix(".schema.strict.json")?;
    let start = stem.find(|c: char| c.is_uppercase())?;
    let type_name = &stem[start..];
    Some(format!("#/components/schemas/{type_name}"))
}

// ── Pass 1 ────────────────────────────────────────────────────────────────────

fn enrich_from_jsonschema(
    spec: &mut YamlValue,
    jsonschema_dir: &Path,
    camel_case: bool,
    schema_overrides: &HashSet<String>,
) -> Result<()> {
    let pattern = jsonschema_dir
        .join("*.schema.strict.bundle.json")
        .to_string_lossy()
        .into_owned();

    let mut files: Vec<std::path::PathBuf> = glob(&pattern)
        .map_err(|e| crate::Error::Build(format!("Glob pattern error: {e}")))?
        .filter_map(|r: Result<std::path::PathBuf, _>| r.ok())
        .collect();
    files.sort();

    if files.is_empty() {
        eprintln!(
            "enrich-openapi: no JSON Schema files found in {}",
            jsonschema_dir.display()
        );
        return Ok(());
    }

    let mut updated = 0usize;
    let mut added = 0usize;

    for path in &files {
        let Some(SchemaBundle {
            type_name,
            defs,
            root_schema,
        }) = load_bundle(path)?
        else {
            continue;
        };

        // Navigate to components.schemas.<TypeName>
        let schemas = spec
            .get_mut("components")
            .and_then(|c| c.get_mut("schemas"));
        let schemas = match schemas {
            Some(s) => s,
            None => {
                eprintln!("enrich-openapi: openapi.yaml has no components.schemas, skipping");
                break;
            }
        };

        if insert_missing_schema(schemas, &type_name, &root_schema) {
            added += 1;
        } else {
            updated += 1;
        }

        let oa_schema = schemas
            .as_mapping_mut()
            .and_then(|m| m.get_mut(type_name.as_str()));
        if let Some(oa) = oa_schema {
            // Authoritative override: gnostic can clobber a proto message whose name
            // collides with one of its own OpenAPI meta-model types (e.g. a
            // `example.schemas.v1.Schema` is overwritten by gnostic's `Schema`),
            // leaving the emitted component with meta-object fields instead of the
            // real entity's. For such keys, rebuild the component body from the
            // authoritative JSON Schema root before the usual validation merge.
            if schema_overrides.contains(type_name.as_str()) {
                *oa = json_to_yaml(&root_schema);
                // Drop JSON-Schema-only keys that must not leak into the OpenAPI spec.
                if let Some(m) = oa.as_mapping_mut() {
                    m.remove(YamlValue::String("$schema".into()));
                }
            }
            enrich_schema(oa, &root_schema, &defs, camel_case);
        }
    }

    println!(
        "enrich-openapi: enriched {} schemas ({} updated, {} added)",
        updated + added,
        updated,
        added
    );
    Ok(())
}

/// A parsed `*.schema.strict.bundle.json` file.
struct SchemaBundle {
    /// The OpenAPI schema key (the proto message name).
    type_name: String,
    /// The bundle's `$defs`.
    defs: JsonMap<String, serde_json::Value>,
    /// The `$defs` entry the bundle's root `$ref` points at.
    root_schema: serde_json::Value,
}

/// Read and parse a JSON Schema bundle, or `None` (logged) when it cannot be used.
fn load_bundle(path: &Path) -> Result<Option<SchemaBundle>> {
    let filename = path
        .file_name()
        .and_then(|f: &std::ffi::OsStr| f.to_str())
        .unwrap_or_default();
    let Some(type_name) = type_name_from_filename(filename) else {
        eprintln!("enrich-openapi: skipping {filename} (cannot extract type name)");
        return Ok(None);
    };

    let content = fs::read_to_string(path).map_err(|e| file_error("read", path, e))?;
    let bundle: serde_json::Value =
        serde_json::from_str(&content).map_err(|e| file_error("parse JSON", path, e))?;

    let defs: JsonMap<String, serde_json::Value> = bundle
        .get("$defs")
        .and_then(|v: &serde_json::Value| v.as_object())
        .cloned()
        .unwrap_or_default();

    let root_ref = bundle
        .get("$ref")
        .and_then(|v: &serde_json::Value| v.as_str())
        .unwrap_or("");
    // Surface an unexpected `$ref` shape instead of silently coercing it to an empty key,
    // which would mask malformed bundles behind a generic "could not resolve" message.
    let Some(root_key) = root_ref.strip_prefix("#/$defs/") else {
        eprintln!(
            "enrich-openapi: unexpected root $ref format '{root_ref}' for {type_name} \
             (expected '#/$defs/<name>'), skipping"
        );
        return Ok(None);
    };
    let Some(root_schema) = defs.get(root_key).cloned() else {
        eprintln!(
            "enrich-openapi: could not resolve root $ref '{root_ref}' for {type_name}, skipping"
        );
        return Ok(None);
    };

    Ok(Some(SchemaBundle {
        type_name,
        defs,
        root_schema,
    }))
}

/// Add a `{type: <root type>}` stub for `type_name` when `schemas` lacks it; returns whether it
/// was added.
fn insert_missing_schema(
    schemas: &mut YamlValue,
    type_name: &str,
    root_schema: &serde_json::Value,
) -> bool {
    let exists = schemas
        .as_mapping()
        .map(|m| m.contains_key(type_name))
        .unwrap_or(false);
    if exists {
        return false;
    }

    let ty = root_schema
        .get("type")
        .and_then(|v: &serde_json::Value| v.as_str())
        .unwrap_or("object");
    if let Some(map) = schemas.as_mapping_mut() {
        let mut entry = serde_yaml::Mapping::new();
        entry.insert(
            YamlValue::String("type".into()),
            YamlValue::String(ty.to_string()),
        );
        map.insert(
            YamlValue::String(type_name.to_string()),
            YamlValue::Mapping(entry),
        );
    }
    true
}

/// Extract type name from a JSON Schema bundle filename.
/// e.g. "dda.coordinator.v1.CreateWaveRequest.schema.strict.bundle.json" → "CreateWaveRequest"
fn type_name_from_filename(filename: &str) -> Option<String> {
    let stem = filename.strip_suffix(".schema.strict.bundle.json")?;
    let start = stem.find(|c: char| c.is_uppercase())?;
    Some(stem[start..].to_string())
}

/// Resolve a `#/$defs/<key>` reference within a bundle's $defs map.
fn resolve_ref<'a>(
    ref_str: &str,
    defs: &'a JsonMap<String, serde_json::Value>,
) -> Option<&'a serde_json::Value> {
    let key = ref_str.strip_prefix("#/$defs/")?;
    defs.get(key)
}

/// The `$defs` entry `node`'s `$ref` points at, or `node` itself when it has no resolvable
/// `$ref`.
fn resolve_or_self<'a>(
    node: &'a serde_json::Value,
    defs: &'a JsonMap<String, serde_json::Value>,
) -> &'a serde_json::Value {
    node.get("$ref")
        .and_then(|v| v.as_str())
        .and_then(|ref_str| resolve_ref(ref_str, defs))
        .unwrap_or(node)
}

const VALIDATION_FIELDS: &[&str] = &[
    "minLength",
    "maxLength",
    "pattern",
    "minimum",
    "maximum",
    "exclusiveMinimum",
    "exclusiveMaximum",
    "minItems",
    "maxItems",
    "enum",
    "additionalProperties",
    "required",
    "description",
    "title",
];

/// Merge validation fields from a JSON Schema node into an OpenAPI YAML schema node.
fn merge_validation(source: &serde_json::Value, target: &mut YamlValue) {
    for &key in VALIDATION_FIELDS {
        let val = match source.get(key) {
            Some(v) => v,
            None => continue,
        };

        // JSON Schema 2020-12 numeric form → OpenAPI 3.0 boolean form.
        let bound = match key {
            "exclusiveMinimum" => Some("minimum"),
            "exclusiveMaximum" => Some("maximum"),
            _ => None,
        };
        if let Some(bound) = bound {
            if let Some(n) = val.as_f64() {
                yaml_set(
                    target,
                    bound,
                    YamlValue::Number(serde_yaml::Number::from(n)),
                );
                yaml_set(target, key, YamlValue::Bool(true));
            }
            continue;
        }

        yaml_set(target, key, json_to_yaml(val));
    }
}

/// Recursively enrich an OpenAPI schema YAML node with validation from a JSON Schema node.
fn enrich_schema(
    openapi: &mut YamlValue,
    json_schema: &serde_json::Value,
    defs: &JsonMap<String, serde_json::Value>,
    camel_case: bool,
) {
    merge_validation(json_schema, openapi);

    // Recurse into properties
    if let Some(js_props) = json_schema.get("properties").and_then(|v| v.as_object()) {
        for (snake_key, js_prop) in js_props {
            let lookup_key = if camel_case {
                crate::utils::strings::lower_camel(snake_key)
            } else {
                snake_key.clone()
            };

            if let Some(oa_prop) = openapi
                .get_mut("properties")
                .and_then(|p| p.get_mut(lookup_key.as_str()))
            {
                enrich_schema(oa_prop, resolve_or_self(js_prop, defs), defs, camel_case);
            }
        }
    }

    // Recurse into items
    if let Some(js_items) = json_schema.get("items")
        && let Some(oa_items) = openapi.get_mut("items")
    {
        enrich_schema(oa_items, resolve_or_self(js_items, defs), defs, camel_case);
    }

    // Recurse into combiners
    for combiner in &["allOf", "oneOf", "anyOf"] {
        if let Some(js_list) = json_schema.get(combiner).and_then(|v| v.as_array())
            && let Some(oa_list) = openapi.get_mut(combiner).and_then(|v| v.as_sequence_mut())
        {
            for (oa_entry, js_entry) in oa_list.iter_mut().zip(js_list) {
                enrich_schema(oa_entry, resolve_or_self(js_entry, defs), defs, camel_case);
            }
        }
    }
}

// ── Pass 2 ────────────────────────────────────────────────────────────────────

fn dedup_path_params(spec: &mut YamlValue, metadata: &CodeGenMetadata) {
    let mut removed_total = 0usize;

    for service in metadata.services.values() {
        for method in &service.methods {
            let path_params: HashSet<&str> = method
                .http_pattern
                .parameters
                .iter()
                .map(String::as_str)
                .collect();
            if path_params.is_empty() {
                continue;
            }

            let input_type = last_segment(&method.input_type);

            let schema = spec
                .get_mut("components")
                .and_then(|c| c.get_mut("schemas"))
                .and_then(|s| s.get_mut(input_type));

            let schema = match schema {
                Some(s) => s,
                None => continue,
            };

            // Remove from properties
            if let Some(props) = schema
                .get_mut("properties")
                .and_then(|p| p.as_mapping_mut())
            {
                let before = props.len();
                props.retain(|k, _| {
                    let key = k.as_str().unwrap_or("");
                    !path_params.contains(key)
                });
                removed_total += before - props.len();
            }

            // Remove from required array
            if let Some(required) = schema.get_mut("required").and_then(|r| r.as_sequence_mut()) {
                required.retain(|v| {
                    let key = v.as_str().unwrap_or("");
                    !path_params.contains(key)
                });
            }
        }
    }

    if removed_total > 0 {
        println!(
            "enrich-openapi: dedup removed {removed_total} path-bound field(s) from request body schemas"
        );
    }
}

// ── Pass 3 ──────────────────────────────────────────────────────────────────

/// Rename `components/schemas` map keys and rewrite every `$ref` that points at a
/// renamed key.
///
/// The OpenAPI schema key is derived by gnostic from the proto message name, and
/// gnostic offers no per-message key override. This pass decouples the two: the
/// proto message (and thus every Rust/Python/Node model type) keeps its clean
/// name, while the emitted spec presents a different key — e.g. to align the spec
/// with a reference spec's naming (`Catalog` -> `CatalogInfo`) so downstream OpenAPI
/// consumers match it without churning the internal model API.
///
/// `renames` maps current schema key -> desired schema key. Entries whose source
/// key is absent from the spec are ignored (logged), so a rename table can name
/// keys that only some builds emit.
fn rename_schemas(spec: &mut YamlValue, renames: &BTreeMap<String, String>) {
    if renames.is_empty() {
        return;
    }

    // Rename the component definitions themselves. Preserve map order by rebuilding
    // the schemas mapping with keys substituted in place.
    let mut applied = 0usize;
    if let Some(schemas) = spec
        .get_mut("components")
        .and_then(|c| c.get_mut("schemas"))
        .and_then(|s| s.as_mapping_mut())
    {
        let mut rebuilt = serde_yaml::Mapping::new();
        for (k, v) in std::mem::take(schemas) {
            match k.as_str().and_then(|key| renames.get(key)) {
                Some(new_key) => {
                    rebuilt.insert(YamlValue::String(new_key.clone()), v);
                    applied += 1;
                }
                None => {
                    rebuilt.insert(k, v);
                }
            }
        }
        *schemas = rebuilt;
    }

    // Rewrite every `$ref: #/components/schemas/<old>` to the new key, everywhere.
    rewrite_schema_refs(spec, renames);

    let requested = renames.len();
    if applied < requested {
        eprintln!(
            "enrich-openapi: schema rename applied {applied}/{requested} (some source keys were absent from the spec)"
        );
    } else {
        eprintln!("enrich-openapi: schema rename applied {applied} key(s)");
    }
}

fn rewrite_schema_refs(node: &mut YamlValue, renames: &BTreeMap<String, String>) {
    match node {
        YamlValue::Mapping(map) => {
            if let Some(YamlValue::String(s)) = map.get_mut("$ref")
                && let Some(old) = s.strip_prefix("#/components/schemas/")
                && let Some(new_key) = renames.get(old)
            {
                *s = format!("#/components/schemas/{new_key}");
            }
            for (_, v) in map.iter_mut() {
                rewrite_schema_refs(v, renames);
            }
        }
        YamlValue::Sequence(seq) => {
            for item in seq.iter_mut() {
                rewrite_schema_refs(item, renames);
            }
        }
        _ => {}
    }
}

// ── YAML helpers ──────────────────────────────────────────────────────────────

fn yaml_set(target: &mut YamlValue, key: &str, value: YamlValue) {
    if let YamlValue::Mapping(m) = target {
        m.insert(YamlValue::String(key.to_string()), value);
    }
}

fn json_to_yaml(v: &serde_json::Value) -> YamlValue {
    match v {
        serde_json::Value::Null => YamlValue::Null,
        serde_json::Value::Bool(b) => YamlValue::Bool(*b),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                YamlValue::Number(serde_yaml::Number::from(i))
            } else if let Some(f) = n.as_f64() {
                YamlValue::Number(serde_yaml::Number::from(f))
            } else {
                YamlValue::String(n.to_string())
            }
        }
        serde_json::Value::String(s) => YamlValue::String(s.clone()),
        serde_json::Value::Array(arr) => {
            YamlValue::Sequence(arr.iter().map(json_to_yaml).collect())
        }
        serde_json::Value::Object(obj) => {
            let mut mapping = serde_yaml::Mapping::new();
            for (k, val) in obj {
                mapping.insert(YamlValue::String(k.clone()), json_to_yaml(val));
            }
            YamlValue::Mapping(mapping)
        }
    }
}