use pgrx::*;
use jsonschema::{Draft, Validator};
use serde_json::json;
use std::collections::HashMap;
use std::sync::RwLock;
use lazy_static::lazy_static;
use jsonschema;

pg_module_magic!();

// Global, thread-safe schema cache using the correct Validator type
lazy_static! {
    static ref SCHEMA_CACHE: RwLock<HashMap<String, Validator>> = RwLock::new(HashMap::new());
}

// Cache a schema explicitly with a provided ID
#[pg_extern(immutable, strict, parallel_safe)]
fn cache_schema(schema_id: &str, schema: JsonB) -> bool {
    match jsonschema::options()
        .with_draft(Draft::Draft7)
        .should_validate_formats(true)
        .build(&schema.0)
    {
        Ok(compiled) => {
            SCHEMA_CACHE.write().unwrap().insert(schema_id.to_string(), compiled);
            true
        },
        Err(e) => {
            notice!("Failed to cache schema '{}': {}", schema_id, e);
            false
        }
    }
}

// Check if a schema is cached
#[pg_extern(immutable, strict, parallel_safe)]
fn schema_cached(schema_id: &str) -> bool {
    SCHEMA_CACHE.read().unwrap().contains_key(schema_id)
}

// Validate JSONB instance against a cached schema by ID
#[pg_extern(immutable, strict, parallel_safe)]
fn validate_schema(schema_id: &str, instance: JsonB) -> JsonB {
    let cache = SCHEMA_CACHE.read().unwrap();
    let compiled_schema: &Validator = match cache.get(schema_id) {
        Some(schema) => schema,
        None => {
            return JsonB(json!({
                "valid": false,
                "errors": [format!("Schema ID '{}' not cached", schema_id)]
            }));
        }
    };

    if compiled_schema.is_valid(&instance.0) {
        JsonB(json!({ "valid": true }))
    } else {
        let errors: Vec<String> = compiled_schema
            .iter_errors(&instance.0)
            .map(|e| e.to_string())
            .collect();

        JsonB(json!({ "valid": false, "errors": errors }))
    }
}

// Clear the entire schema cache explicitly
#[pg_extern(immutable, parallel_safe)]
fn clear_schema_cache() -> bool {
    SCHEMA_CACHE.write().unwrap().clear();
    true
}

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

    #[pg_test]
    fn test_cache_and_validate_schema() {
        assert!(crate::cache_schema("test_schema", JsonB(json!({ "type": "object" }))));
        assert!(crate::schema_cached("test_schema"));

        let result_valid = crate::validate_schema("test_schema", JsonB(json!({ "foo": "bar" })));
        assert_eq!(result_valid.0["valid"], true);

        let result_invalid = crate::validate_schema("test_schema", JsonB(json!(42)));
        assert_eq!(result_invalid.0["valid"], false);
        assert!(result_invalid.0["errors"][0].as_str().unwrap().contains("not of type"));
    }

    #[pg_test]
    fn test_schema_not_cached() {
        let result = crate::validate_schema("unknown_schema", JsonB(json!({})));
        assert_eq!(result.0["valid"], false);
        assert!(result.0["errors"][0].as_str().unwrap().contains("not cached"));
    }

    #[pg_test]
    fn test_clear_schema_cache() {
        crate::cache_schema("clear_test", JsonB(json!({ "type": "object" })));
        assert!(crate::schema_cached("clear_test"));

        crate::clear_schema_cache();
        assert!(!crate::schema_cached("clear_test"));
    }

    #[pg_test]
    fn test_invalid_schema_cache() {
        let result = crate::cache_schema("bad_schema", JsonB(json!({ "type": "unknown_type" })));
        assert!(!result);
        assert!(!crate::schema_cached("bad_schema"));
    }
}

#[cfg(test)]
pub mod pg_test {
    pub fn setup(_options: Vec<&str>) {
        // Initialization if needed
    }

    pub fn postgresql_conf_options() -> Vec<&'static str> {
        vec![]
    }
}
