pmcp 2.8.1

//! Schema utilities for normalizing and inlining JSON schemas
//!
//! Provides utilities to process JSON schemas generated by schemars,
//! including inlining $ref references for better client compatibility.

#[cfg(feature = "schema-generation")]
use serde_json::{json, Value};

/// Configuration for schema normalization
#[cfg(feature = "schema-generation")]
#[derive(Debug, Clone)]
pub struct NormalizerConfig {
    /// Maximum depth for inlining references (default: 10)
    pub max_inline_depth: usize,
    /// Maximum size in bytes for inlined schemas (default: 100KB)
    pub max_inline_size: usize,
    /// Whether to keep definitions section (default: false)
    pub keep_definitions: bool,
    /// Whether to remove schema metadata like $schema, $id (default: true)
    pub remove_metadata: bool,
}

impl Default for NormalizerConfig {
    fn default() -> Self {
        Self {
            max_inline_depth: 10,
            max_inline_size: 100_000, // 100KB
            keep_definitions: false,
            remove_metadata: true,
        }
    }
}

/// Normalize a JSON schema by inlining all $ref references
///
/// This function takes a schema generated by schemars and processes it to:
/// - Inline all $ref references (up to configured depth)
/// - Optionally remove the definitions section
/// - Simplify the schema structure
///
/// # Example
/// ```no_run
/// use schemars::schema_for;
/// use pmcp::server::schema_utils::normalize_schema;
///
/// #[derive(schemars::JsonSchema)]
/// struct MyType {
///     field: String,
/// }
///
/// let schema = schema_for!(MyType);
/// let normalized = normalize_schema(serde_json::to_value(schema).unwrap());
/// ```
#[cfg(feature = "schema-generation")]
pub fn normalize_schema(schema: Value) -> Value {
    normalize_schema_with_config(schema, &NormalizerConfig::default())
}

/// Strip schema metadata fields (`$schema`, `$id`) in place.
#[cfg(feature = "schema-generation")]
fn strip_schema_metadata(schema: &mut Value) {
    let Some(obj) = schema.as_object_mut() else {
        return;
    };
    obj.remove("$schema");
    obj.remove("$id");
}

/// Extract (removing from `schema` unless `keep_definitions` is true) the
/// `definitions` / `$defs` block, returning it as an owned `Value` if present.
#[cfg(feature = "schema-generation")]
fn extract_definitions_block(schema: &mut Value, keep_definitions: bool) -> Option<Value> {
    if keep_definitions {
        return schema
            .as_object()
            .and_then(|obj| obj.get("definitions").or_else(|| obj.get("$defs")))
            .cloned();
    }
    schema
        .as_object_mut()
        .and_then(|obj| obj.remove("definitions"))
        .or_else(|| schema.as_object_mut().and_then(|obj| obj.remove("$defs")))
}

/// Attempt to inline a root-level `$ref` pointing into `#/definitions/` or
/// `#/$defs/`, returning the replacement value if found.
#[cfg(feature = "schema-generation")]
fn try_inline_root_ref(schema: &Value, defs: Option<&Value>) -> Option<Value> {
    let obj = schema.as_object()?;
    let Value::String(ref_str) = obj.get("$ref")? else {
        return None;
    };
    if !ref_str.starts_with("#/definitions/") && !ref_str.starts_with("#/$defs/") {
        return None;
    }
    let def_name = ref_str.rsplit('/').next().unwrap_or("");
    defs?.get(def_name).cloned()
}

/// Apply `extract_definitions_block` output: inline refs and restore the
/// definitions block if requested.
#[cfg(feature = "schema-generation")]
fn apply_inline_pass(
    schema: &mut Value,
    definitions: Value,
    schema_size: usize,
    config: &NormalizerConfig,
) {
    let mut context = InlineContext {
        definitions: &definitions,
        current_depth: 0,
        max_depth: config.max_inline_depth,
        current_size: schema_size,
        max_size: config.max_inline_size,
    };
    inline_refs_with_context(schema, &mut context);
    if config.keep_definitions {
        if let Some(obj) = schema.as_object_mut() {
            obj.insert("definitions".to_string(), definitions);
        }
    }
}

/// Normalize a JSON schema with custom configuration.
///
/// Refactored in 75-01 Task 1a-B (P1): extracted
/// `strip_schema_metadata`, `extract_definitions_block`,
/// `try_inline_root_ref`, `apply_inline_pass` so this orchestrator is a
/// sequential pipeline.
#[cfg(feature = "schema-generation")]
pub fn normalize_schema_with_config(mut schema: Value, config: &NormalizerConfig) -> Value {
    let schema_size = serde_json::to_string(&schema).unwrap_or_default().len();
    if schema_size > config.max_inline_size {
        // Schema too large, return with minimal processing
        if config.remove_metadata {
            strip_schema_metadata(&mut schema);
        }
        return schema;
    }

    let definitions = extract_definitions_block(&mut schema, config.keep_definitions);
    let defs_clone = definitions.clone();

    if let Some(defs) = definitions {
        apply_inline_pass(&mut schema, defs, schema_size, config);
    }

    if !config.remove_metadata {
        return schema;
    }
    strip_schema_metadata(&mut schema);

    if !config.keep_definitions {
        if let Some(inlined) = try_inline_root_ref(&schema, defs_clone.as_ref()) {
            return inlined;
        }
    }

    schema
}

/// Context for tracking inlining depth and size
#[cfg(feature = "schema-generation")]
struct InlineContext<'a> {
    definitions: &'a Value,
    current_depth: usize,
    max_depth: usize,
    current_size: usize,
    max_size: usize,
}

/// Parse a `$ref` string pointing at `#/definitions/<name>` or
/// `#/$defs/<name>`, returning the definition name.
#[cfg(feature = "schema-generation")]
fn parse_local_ref_name(ref_str: &str) -> Option<&str> {
    if ref_str.starts_with("#/definitions/") || ref_str.starts_with("#/$defs/") {
        return ref_str.rsplit('/').next();
    }
    None
}

/// Attempt to inline a single `$ref` in `map`, subject to the context's
/// size budget. Returns `true` if the ref was inlined.
#[cfg(feature = "schema-generation")]
fn try_inline_single_ref(
    map: &mut serde_json::Map<String, Value>,
    context: &InlineContext<'_>,
) -> bool {
    let Some(Value::String(ref_str)) = map.get("$ref") else {
        return false;
    };
    let Some(def_name) = parse_local_ref_name(ref_str) else {
        return false;
    };
    let Some(def_value) = context.definitions.get(def_name) else {
        return false;
    };
    let def_size = serde_json::to_string(&def_value).unwrap_or_default().len();
    if context.current_size + def_size > context.max_size {
        return false;
    }
    let Some(def_obj) = def_value.as_object() else {
        return false;
    };
    map.remove("$ref");
    for (key, val) in def_obj {
        if !map.contains_key(key) {
            map.insert(key.clone(), val.clone());
        }
    }
    true
}

/// Walk an object, inline any resolvable `$ref`, and recurse into children.
#[cfg(feature = "schema-generation")]
fn inline_refs_in_object(
    map: &mut serde_json::Map<String, Value>,
    context: &mut InlineContext<'_>,
) {
    let _ = try_inline_single_ref(map, context);
    for val in map.values_mut() {
        inline_refs_with_context(val, context);
    }
}

/// Recursively inline references with depth and size tracking.
///
/// Refactored in 75-01 Task 1a-B (P1): extracted [`parse_local_ref_name`],
/// [`try_inline_single_ref`], [`inline_refs_in_object`] so this function
/// reads as a flat depth/size-guarded recursion.
#[cfg(feature = "schema-generation")]
fn inline_refs_with_context(value: &mut Value, context: &mut InlineContext<'_>) {
    if context.current_depth >= context.max_depth {
        return;
    }
    let value_size = serde_json::to_string(&value).unwrap_or_default().len();
    if context.current_size + value_size > context.max_size {
        return;
    }
    context.current_depth += 1;
    context.current_size += value_size;

    match value {
        Value::Object(map) => inline_refs_in_object(map, context),
        Value::Array(arr) => {
            for item in arr.iter_mut() {
                inline_refs_with_context(item, context);
            }
        },
        _ => {},
    }

    context.current_depth -= 1;
}

// Note: the legacy `inline_refs` helper was removed in 75-01 Task 1a-B as
// dead code — confirmed no callers via workspace grep. The current
// implementation is `inline_refs_with_context` with depth/size tracking.

/// Create a simple schema for primitive types
///
/// Helper function to create schemas for basic types without using schemars
#[cfg(feature = "schema-generation")]
pub fn simple_schema(type_name: &str, description: Option<&str>) -> Value {
    let mut schema = json!({
        "type": match type_name {
            "string" | "String" => "string",
            "number" | "f64" | "f32" => "number",
            "integer" | "i32" | "i64" | "u32" | "u64" | "usize" | "isize" => "integer",
            "boolean" | "bool" => "boolean",
            "array" | "Vec" => "array",
            "object" | "HashMap" | "BTreeMap" => "object",
            _ => "string", // Default to string
        }
    });

    if let Some(desc) = description {
        schema["description"] = json!(desc);
    }

    schema
}

/// Add validation constraints to a schema
///
/// Adds common validation constraints to a schema object
#[cfg(feature = "schema-generation")]
pub fn add_constraints(mut schema: Value, constraints: SchemaConstraints) -> Value {
    if let Some(obj) = schema.as_object_mut() {
        // String constraints
        if let Some(min_length) = constraints.min_length {
            obj.insert("minLength".to_string(), json!(min_length));
        }
        if let Some(max_length) = constraints.max_length {
            obj.insert("maxLength".to_string(), json!(max_length));
        }
        if let Some(pattern) = constraints.pattern {
            obj.insert("pattern".to_string(), json!(pattern));
        }

        // Numeric constraints
        if let Some(minimum) = constraints.minimum {
            obj.insert("minimum".to_string(), json!(minimum));
        }
        if let Some(maximum) = constraints.maximum {
            obj.insert("maximum".to_string(), json!(maximum));
        }
        if let Some(exclusive_minimum) = constraints.exclusive_minimum {
            obj.insert("exclusiveMinimum".to_string(), json!(exclusive_minimum));
        }
        if let Some(exclusive_maximum) = constraints.exclusive_maximum {
            obj.insert("exclusiveMaximum".to_string(), json!(exclusive_maximum));
        }

        // Array constraints
        if let Some(min_items) = constraints.min_items {
            obj.insert("minItems".to_string(), json!(min_items));
        }
        if let Some(max_items) = constraints.max_items {
            obj.insert("maxItems".to_string(), json!(max_items));
        }
        if let Some(unique_items) = constraints.unique_items {
            obj.insert("uniqueItems".to_string(), json!(unique_items));
        }

        // Enum constraint
        if let Some(enum_values) = constraints.enum_values {
            obj.insert("enum".to_string(), json!(enum_values));
        }
    }

    schema
}

/// Schema constraints that can be added to a schema
#[cfg(feature = "schema-generation")]
#[derive(Debug, Default)]
pub struct SchemaConstraints {
    /// Minimum length for string values
    pub min_length: Option<usize>,
    /// Maximum length for string values
    pub max_length: Option<usize>,
    /// Regex pattern that string values must match
    pub pattern: Option<String>,

    /// Minimum value for numeric types
    pub minimum: Option<f64>,
    /// Maximum value for numeric types
    pub maximum: Option<f64>,
    /// Exclusive minimum value for numeric types
    pub exclusive_minimum: Option<f64>,
    /// Exclusive maximum value for numeric types
    pub exclusive_maximum: Option<f64>,

    /// Minimum number of items in arrays
    pub min_items: Option<usize>,
    /// Maximum number of items in arrays
    pub max_items: Option<usize>,
    /// Whether array items must be unique
    pub unique_items: Option<bool>,

    /// List of allowed values (enum constraint)
    pub enum_values: Option<Vec<Value>>,
}

#[cfg(all(test, feature = "schema-generation"))]
mod tests {
    use super::*;
    use schemars::JsonSchema;
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Deserialize, Serialize, JsonSchema)]
    struct TestStruct {
        /// A simple string field
        name: String,
        /// An optional number
        age: Option<u32>,
        /// A nested structure
        address: Address,
    }

    #[derive(Debug, Deserialize, Serialize, JsonSchema)]
    struct Address {
        street: String,
        city: String,
    }

    #[test]
    fn test_normalize_schema() {
        let schema = schemars::schema_for!(TestStruct);
        let json_schema = serde_json::to_value(schema).unwrap();
        let normalized = normalize_schema(json_schema);

        // Check that $schema and definitions are removed
        assert!(!normalized.as_object().unwrap().contains_key("$schema"));
        assert!(!normalized.as_object().unwrap().contains_key("definitions"));

        // Check that the properties are present
        let props = normalized["properties"].as_object().unwrap();
        assert!(props.contains_key("name"));
        assert!(props.contains_key("age"));
        assert!(props.contains_key("address"));

        // Check that the address is inlined
        let address_props = props["address"]["properties"].as_object().unwrap();
        assert!(address_props.contains_key("street"));
        assert!(address_props.contains_key("city"));
    }

    #[test]
    fn test_simple_schema() {
        let string_schema = simple_schema("string", Some("A test string"));
        assert_eq!(string_schema["type"], "string");
        assert_eq!(string_schema["description"], "A test string");

        let number_schema = simple_schema("f64", None);
        assert_eq!(number_schema["type"], "number");

        let int_schema = simple_schema("i32", None);
        assert_eq!(int_schema["type"], "integer");
    }

    #[test]
    fn test_add_constraints() {
        let base_schema = json!({"type": "string"});
        let constraints = SchemaConstraints {
            min_length: Some(5),
            max_length: Some(20),
            pattern: Some(r"^\w+$".to_string()),
            ..Default::default()
        };

        let constrained = add_constraints(base_schema, constraints);
        assert_eq!(constrained["minLength"], 5);
        assert_eq!(constrained["maxLength"], 20);
        assert_eq!(constrained["pattern"], r"^\w+$");
    }
}