jspg/src/lib.rs

use pgrx::*;

pg_module_magic!();

use serde_json::{json, Value};
use std::{collections::HashMap, sync::RwLock};
use boon::{Compiler, Schemas, ValidationError, SchemaIndex};
use lazy_static::lazy_static;

struct BoonCache {
    schemas: Schemas,
    id_to_index: HashMap<String, SchemaIndex>,
}

lazy_static! {
    static ref SCHEMA_CACHE: RwLock<BoonCache> = RwLock::new(BoonCache {
        schemas: Schemas::new(),
        id_to_index: HashMap::new()
    });
}

#[pg_extern(strict)]
fn cache_json_schema(schema_id: &str, schema: JsonB) -> JsonB {
    let mut cache = SCHEMA_CACHE.write().unwrap();
    let schema_value: Value = schema.0;

    let mut compiler = Compiler::new();
    compiler.enable_format_assertions();

    let schema_url = format!("urn:jspg:{}", schema_id);

    if let Err(e) = compiler.add_resource(&schema_url, schema_value) {
        return JsonB(json!({
            "success": false,
            "error": format!("Failed to add schema resource '{}': {}", schema_id, e)
        }));
    }

    match compiler.compile(&schema_url, &mut cache.schemas) {
        Ok(sch_index) => {
            cache.id_to_index.insert(schema_id.to_string(), sch_index);
            JsonB(json!({
                "success": true,
                "schema_id": schema_id,
                "message": "Schema cached successfully."
            }))
        }
        Err(e) => JsonB(json!({
            "success": false,
            "schema_id": schema_id,
            "error": format!("Schema compilation failed: {}", e)
        })),
    }
}

#[pg_extern(strict, parallel_safe)]
fn validate_json_schema(schema_id: &str, instance: JsonB) -> JsonB {
    let cache = SCHEMA_CACHE.read().unwrap();

    match cache.id_to_index.get(schema_id) {
        None => JsonB(json!({
            "success": false,
            "errors": [{
                "kind": "SchemaNotFound",
                "message": format!("Schema with id '{}' not found in cache", schema_id)
            }]
        })),
        Some(sch_index) => {
            let instance_value: Value = instance.0;
            match cache.schemas.validate(&instance_value, *sch_index) {
                Ok(_) => JsonB(json!({ "success": true })),
                Err(validation_error) => {
                    let error_details = format_boon_errors(&validation_error);
                    JsonB(json!({
                        "success": false,
                        "errors": [error_details]
                    }))
                }
            }
        }
    }
}

fn format_boon_errors(error: &ValidationError) -> Value {
    json!({
        "instance_path": error.instance_location.to_string(),
        "schema_path": error.schema_url.to_string(),
        "kind": format!("{:?}", error.kind),
        "message": format!("{}", error),
        "causes": error
            .causes
            .iter()
            .map(format_boon_errors)
            .collect::<Vec<_>>()
    })
}

#[pg_extern(strict, parallel_safe)]
fn json_schema_cached(schema_id: &str) -> bool {
    let cache = SCHEMA_CACHE.read().unwrap();
    cache.id_to_index.contains_key(schema_id)
}

#[pg_extern(strict)]
fn clear_json_schemas() -> JsonB {
    let mut cache = SCHEMA_CACHE.write().unwrap();
    *cache = BoonCache {
        schemas: Schemas::new(),
        id_to_index: HashMap::new()
    };
    JsonB(json!({
        "success": true,
        "message": "Schema cache cleared."
    }))
}

#[pg_extern(strict, parallel_safe)]
fn show_json_schemas() -> JsonB {
    let cache = SCHEMA_CACHE.read().unwrap();
    let ids: Vec<&String> = cache.id_to_index.keys().collect();
    JsonB(json!({
        "cached_schema_ids": ids
    }))
}

#[pg_schema]
#[cfg(any(test, feature = "pg_test"))]
mod tests {
    use pgrx::*;
    use pgrx::pg_test;
    use super::*;
    use serde_json::json;

    fn jsonb(val: Value) -> JsonB {
        JsonB(val)
    }

    fn setup_test() {
        clear_json_schemas();
    }

    #[pg_test]
    fn test_cache_and_validate_json_schema() {
        setup_test();
        let schema_id = "my_schema";
        let schema = json!({
            "type": "object",
            "properties": {
                "name": { "type": "string" },
                "age": { "type": "integer", "minimum": 0 }
            },
            "required": ["name", "age"]
        });
        let valid_instance = json!({ "name": "Alice", "age": 30 });
        let invalid_instance_type = json!({ "name": "Bob", "age": -5 });
        let invalid_instance_missing = json!({ "name": "Charlie" });

        let cache_result = cache_json_schema(schema_id, jsonb(schema.clone()));
        assert!(cache_result.0["success"].as_bool().unwrap());

        let valid_result = validate_json_schema(schema_id, jsonb(valid_instance));
        assert!(valid_result.0["success"].as_bool().unwrap());

        let invalid_result_type = validate_json_schema(schema_id, jsonb(invalid_instance_type));
        assert!(!invalid_result_type.0["success"].as_bool().unwrap());

        // --- Assertions for invalid_result_type --- 

        // Get top-level errors
        let top_level_errors = invalid_result_type.0["errors"].as_array().expect("Top-level 'errors' should be an array");
        assert_eq!(top_level_errors.len(), 1, "Should have exactly one top-level error for invalid type");

        // Get the first (and only) top-level error
        let top_level_error = top_level_errors.get(0).expect("Should get the first top-level error");

        // Check top-level error kind
        assert!(top_level_error.get("kind").and_then(Value::as_str).map_or(false, |k| k.starts_with("Schema { url:")),
                "Incorrect kind for top-level error. Expected 'Schema {{ url:'. Error: {:?}. All errors: {:?}", top_level_error, top_level_errors);

        // Get the 'causes' array from the top-level error
        let causes_age = top_level_error.get("causes").and_then(Value::as_array).expect("Top-level error 'causes' should be an array");
        assert_eq!(causes_age.len(), 1, "Should have one cause for the age error");

        // Get the actual age error from the 'causes' array
        let age_error = causes_age.get(0).expect("Should have an error object in 'causes'");
        assert_eq!(age_error.get("instance_path").and_then(Value::as_str), Some("/age"),
                   "Incorrect instance_path for age error. Error: {:?}. All errors: {:?}", age_error, top_level_errors);

        assert!(age_error.get("kind").and_then(Value::as_str).map_or(false, |k| k.starts_with("Minimum { got:")),
                "Incorrect kind prefix for age error. Expected 'Minimum {{ got:'. Error: {:?}. All errors: {:?}", age_error, top_level_errors);

        let expected_prefix = "at '/age': must be >=0";
        assert!(age_error.get("message")
                    .and_then(Value::as_str)
                    .map_or(false, |m| m.starts_with(expected_prefix)),
                "Incorrect message prefix for age error. Expected prefix '{}'. Error: {:?}. All errors: {:?}",
                expected_prefix, age_error, top_level_errors);

        let invalid_result_missing = validate_json_schema(schema_id, jsonb(invalid_instance_missing));
        assert!(!invalid_result_missing.0["success"].as_bool().unwrap(), "Validation should fail for missing required field");

        // --- Assertions for invalid_result_missing --- 

        // Get top-level errors
        let top_level_errors_missing = invalid_result_missing.0["errors"].as_array().expect("Errors should be an array for missing field");
        assert_eq!(top_level_errors_missing.len(), 1, "Should have one top-level error for missing field");

        // Get the first (and only) top-level error
        let top_error_missing = top_level_errors_missing.get(0).expect("Should get the first top-level missing field error");

        // Check top-level error kind
        assert!(top_error_missing.get("kind").and_then(Value::as_str).map_or(false, |k| k.starts_with("Schema { url:")), 
                "Incorrect kind for missing field top-level error. Error: {:?}. All errors: {:?}", top_error_missing, top_level_errors_missing);

        // Get the 'causes' array from the top-level error
        let causes_missing = top_error_missing.get("causes").and_then(Value::as_array).expect("Causes should be an array for missing field error");
        assert_eq!(causes_missing.len(), 1, "Should have one cause for missing field");

        // Get the actual missing field error from the 'causes' array
        let missing_error = causes_missing.get(0).expect("Should have missing field error object in 'causes'");

        // Assertions on the specific missing field error
        assert_eq!(missing_error.get("instance_path").and_then(Value::as_str), Some(""), 
                   "Incorrect instance_path for missing age error: {:?}", missing_error);
        assert!(missing_error.get("kind").and_then(Value::as_str).map_or(false, |k| k.starts_with("Required { want: [\"age\"]")), 
                "Incorrect kind for missing age error. Expected prefix 'Required {{ want: [\"age\"] }}'. Error: {:?}", missing_error);
    }

    #[pg_test]
    fn test_validate_json_schema_not_cached() {
        setup_test();
        let instance = json!({ "foo": "bar" });
        let result = validate_json_schema("non_existent_schema", jsonb(instance));
        assert!(!result.0["success"].as_bool().unwrap());
        let errors = result.0["errors"].as_array().unwrap();
        assert_eq!(errors.len(), 1);
        assert_eq!(errors[0]["kind"], json!("SchemaNotFound"));
        assert!(errors[0]["message"].as_str().unwrap().contains("non_existent_schema"));
    }

    #[pg_test]
    fn test_cache_invalid_json_schema() {
        setup_test();
        let schema_id = "invalid_schema";
        let invalid_schema_json = "{\"type\": \"string\" \"maxLength\": 5}";
        let invalid_schema_value: Result<Value, _> = serde_json::from_str(invalid_schema_json);
        assert!(invalid_schema_value.is_err(), "Test setup assumes invalid JSON string");

        let schema_representing_invalid = json!({
            "type": 123
        });

        let result = cache_json_schema(schema_id, jsonb(schema_representing_invalid.clone()));
        assert!(!result.0["success"].as_bool().unwrap());
        assert!(result.0["error"].as_str().unwrap().contains("Schema compilation failed"));
    }

    #[pg_test]
    fn test_validate_json_schema_detailed_validation_errors() {
        setup_test();
        let schema_id = "detailed_schema";
        let schema = json!({
            "type": "object",
            "properties": {
                "address": {
                    "type": "object",
                    "properties": {
                        "street": { "type": "string" },
                        "city": { "type": "string", "maxLength": 10 }
                    },
                    "required": ["street", "city"]
                }
            },
            "required": ["address"]
        });
        let invalid_instance = json!({
            "address": {
                "city": "San Francisco Bay Area"
            }
        });

        assert!(cache_json_schema(schema_id, jsonb(schema.clone())).0["success"].as_bool().unwrap());
        let result = validate_json_schema(schema_id, jsonb(invalid_instance));
        assert!(!result.0["success"].as_bool().unwrap());

        let errors = result.0["errors"].as_array().expect("Errors should be an array");
        let top_error = errors.get(0).expect("Expected at least one top-level error object");
        let causes = top_error.get("causes").and_then(Value::as_array).expect("Expected causes array");

        let has_required_street_error = causes.iter().any(|e|
            e.get("instance_path").and_then(Value::as_str) == Some("/address") && // Check path inside cause
            e.get("kind").and_then(Value::as_str).unwrap_or("").starts_with("Required { want:") && // Check kind prefix
            e.get("kind").and_then(Value::as_str).unwrap_or("").contains("street") // Ensure 'street' is mentioned
        );
        assert!(has_required_street_error, "Missing required 'street' error within causes. Actual errors: {:?}", errors);

        let has_maxlength_city_error = causes.iter().any(|e| // Check within causes
            e.get("instance_path").and_then(Value::as_str) == Some("/address/city") &&
            e.get("kind").and_then(Value::as_str).unwrap_or("").starts_with("MaxLength { got:") // Check kind prefix
        );
        assert!(has_maxlength_city_error, "Missing maxLength 'city' error within causes. Actual errors: {:?}", errors);
    }

    #[pg_test]
    fn test_validate_json_schema_oneof_validation_errors() {
        setup_test();
        let schema_id = "oneof_schema";
        let schema = json!({
            "type": "object",
            "properties": {
                "value": {
                    "oneOf": [
                        { "type": "string", "minLength": 5 },
                        { "type": "number", "minimum": 10 }
                    ]
                }
            },
            "required": ["value"]
        });
        assert!(cache_json_schema(schema_id, jsonb(schema.clone())).0["success"].as_bool().unwrap());

        let invalid_instance = json!({ "value": "abc" });
        let result = validate_json_schema(schema_id, jsonb(invalid_instance));

        assert!(!result.0["success"].as_bool().unwrap());

        let errors_val = result.0["errors"].as_array().expect("Errors should be an array");
        let top_schema_error = errors_val.get(0).expect("Expected at least one top-level Schema error object");
        let schema_error_causes = top_schema_error.get("causes").and_then(Value::as_array).expect("Expected causes array for Schema error");

        let oneof_error = schema_error_causes.iter().find(|e| {
            e.get("kind").and_then(Value::as_str) == Some("OneOf(None)") &&
            e.get("instance_path").and_then(Value::as_str) == Some("/value")
        }).expect("Could not find the OneOf(None) error for /value within Schema causes");

        let oneof_causes = oneof_error.get("causes").and_then(Value::as_array)
            .expect("Expected causes array for OneOf error");

        let has_minlength_error = oneof_causes.iter().any(|e| // Check within OneOf causes
            e.get("instance_path").and_then(Value::as_str) == Some("/value") &&
            e.get("kind").and_then(Value::as_str).unwrap_or("").starts_with("MinLength { got:") // Check kind prefix
        );
        assert!(has_minlength_error, "Missing MinLength error within OneOf causes. Actual errors: {:?}", errors_val);

        let has_type_error = oneof_causes.iter().any(|e| // Check within OneOf causes
            e.get("instance_path").and_then(Value::as_str) == Some("/value") &&
            e.get("kind").and_then(Value::as_str).unwrap_or("").starts_with("Type { got: String, want: Types") // More specific kind check
        );
        assert!(has_type_error, "Missing Type error within OneOf causes. Actual errors: {:?}", errors_val);
    }

    #[pg_test]
    fn test_clear_json_schemas() {
        setup_test();
        let schema_id = "schema_to_clear";
        let schema = json!({ "type": "string" });
        cache_json_schema(schema_id, jsonb(schema.clone()));

        let show_result1 = show_json_schemas();
        assert!(show_result1.0["cached_schema_ids"].as_array().unwrap().iter().any(|id| id.as_str() == Some(schema_id)));

        let clear_result = clear_json_schemas();
        assert!(clear_result.0["success"].as_bool().unwrap());

        let show_result2 = show_json_schemas();
        assert!(show_result2.0["cached_schema_ids"].as_array().unwrap().is_empty());

        let instance = json!("test");
        let validate_result = validate_json_schema(schema_id, jsonb(instance));
        assert!(!validate_result.0["success"].as_bool().unwrap());
        assert_eq!(validate_result.0["errors"].as_array().unwrap()[0]["kind"], json!("SchemaNotFound"));
    }

    #[pg_test]
    fn test_show_json_schemas() {
        setup_test();
        let schema_id1 = "schema1";
        let schema_id2 = "schema2";
        let schema = json!({ "type": "boolean" });

        cache_json_schema(schema_id1, jsonb(schema.clone()));
        cache_json_schema(schema_id2, jsonb(schema.clone()));

        let result = show_json_schemas();
        let ids = result.0["cached_schema_ids"].as_array().unwrap();
        assert_eq!(ids.len(), 2);
        assert!(ids.contains(&json!(schema_id1)));
        assert!(ids.contains(&json!(schema_id2)));
    }
}

#[cfg(test)]
pub mod pg_test {
    pub fn setup(_options: Vec<&str>) {
        // perform one-off initialization when the pg_test framework starts
    }

    pub fn postgresql_conf_options() -> Vec<&'static str> {
        // return any postgresql.conf settings that are required for your tests
        vec![]
    }
}