pmcp 2.4.0

//! 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())
}

/// Normalize a JSON schema with custom configuration
#[cfg(feature = "schema-generation")]
pub fn normalize_schema_with_config(mut schema: Value, config: &NormalizerConfig) -> Value {
    // Check size limit early
    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 {
            if let Some(obj) = schema.as_object_mut() {
                obj.remove("$schema");
                obj.remove("$id");
            }
        }
        return schema;
    }

    // Extract definitions if they exist
    let definitions = if config.keep_definitions {
        // Keep definitions but still process refs
        schema
            .as_object()
            .and_then(|obj| obj.get("definitions").or_else(|| obj.get("$defs")))
            .cloned()
    } else {
        // Remove definitions from schema
        schema
            .as_object_mut()
            .and_then(|obj| obj.remove("definitions"))
            .or_else(|| schema.as_object_mut().and_then(|obj| obj.remove("$defs")))
    };

    // Clone definitions for later reference if needed
    let defs_clone = definitions.clone();

    if let Some(defs) = definitions {
        // Inline references with depth tracking
        let mut context = InlineContext {
            definitions: &defs,
            current_depth: 0,
            max_depth: config.max_inline_depth,
            current_size: schema_size,
            max_size: config.max_inline_size,
        };
        inline_refs_with_context(&mut schema, &mut context);

        // If keeping definitions, put them back
        if config.keep_definitions {
            if let Some(obj) = schema.as_object_mut() {
                obj.insert("definitions".to_string(), defs);
            }
        }
    }

    // Remove schema metadata if configured
    if config.remove_metadata {
        if let Some(obj) = schema.as_object_mut() {
            obj.remove("$schema");
            obj.remove("$id");

            // If the schema has a single root definition reference, inline it
            if !config.keep_definitions {
                if let Some(Value::String(ref_str)) = obj.get("$ref") {
                    if ref_str.starts_with("#/definitions/") || ref_str.starts_with("#/$defs/") {
                        let def_name = ref_str.rsplit('/').next().unwrap_or("");
                        if let Some(defs) = defs_clone.as_ref() {
                            if let Some(def_value) = defs.get(def_name) {
                                return def_value.clone();
                            }
                        }
                    }
                }
            }
        }
    }

    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,
}

/// Recursively inline references with depth and size tracking
#[cfg(feature = "schema-generation")]
fn inline_refs_with_context(value: &mut Value, context: &mut InlineContext<'_>) {
    // Check depth limit
    if context.current_depth >= context.max_depth {
        return;
    }

    // Check size limit
    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) => {
            // Check if this object has a $ref
            if let Some(Value::String(ref_str)) = map.get("$ref") {
                if ref_str.starts_with("#/definitions/") || ref_str.starts_with("#/$defs/") {
                    let def_name = ref_str.rsplit('/').next().unwrap_or("");

                    // Find the definition and clone it
                    if let Some(def_value) = context.definitions.get(def_name) {
                        // Check if inlining this would exceed size limit
                        let def_size = serde_json::to_string(&def_value).unwrap_or_default().len();
                        if context.current_size + def_size <= context.max_size {
                            // Replace the entire object with the definition
                            if let Some(def_obj) = def_value.as_object() {
                                map.remove("$ref");
                                for (key, val) in def_obj {
                                    if !map.contains_key(key) {
                                        map.insert(key.clone(), val.clone());
                                    }
                                }
                            }
                        }
                    }
                }
            }

            // Process all values in the object
            for (_key, val) in map.iter_mut() {
                inline_refs_with_context(val, context);
            }
        },
        Value::Array(arr) => {
            // Process all items in the array
            for item in arr.iter_mut() {
                inline_refs_with_context(item, context);
            }
        },
        _ => {},
    }

    context.current_depth -= 1;
}

/// Recursively inline all $ref references in a JSON value (legacy, unlimited)
#[cfg(feature = "schema-generation")]
#[allow(dead_code)]
fn inline_refs(value: &mut Value, definitions: &Value) {
    match value {
        Value::Object(map) => {
            // Check if this object has a $ref
            if let Some(Value::String(ref_str)) = map.get("$ref") {
                if ref_str.starts_with("#/definitions/") || ref_str.starts_with("#/$defs/") {
                    let def_name = ref_str.rsplit('/').next().unwrap_or("");

                    // Find the definition and clone it
                    if let Some(def_value) = definitions.get(def_name) {
                        // Replace the entire object with the definition
                        if let Some(def_obj) = def_value.as_object() {
                            map.remove("$ref");
                            for (key, val) in def_obj {
                                if !map.contains_key(key) {
                                    map.insert(key.clone(), val.clone());
                                }
                            }
                        }
                    }
                }
            }

            // Process all values in the object
            for (_key, val) in map.iter_mut() {
                inline_refs(val, definitions);
            }
        },
        Value::Array(arr) => {
            // Process all items in the array
            for item in arr.iter_mut() {
                inline_refs(item, definitions);
            }
        },
        _ => {},
    }
}

/// 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+$");
    }
}