use pgrx::*;

pg_module_magic!();

// mod schema;
mod registry;
mod validator;
mod util;

use crate::registry::REGISTRY;
// use crate::schema::Schema;
use crate::validator::{Validator, ValidationOptions};
use lazy_static::lazy_static;
use serde_json::{json, Value};
use std::collections::{HashMap, HashSet};

#[derive(Clone, Copy, Debug, PartialEq)]
enum SchemaType {
    Enum,
    Type,
    Family,
    PublicPunc,
    PrivatePunc,
}

struct CachedSchema {
    t: SchemaType,
}

lazy_static! {
    static ref SCHEMA_META: std::sync::RwLock<HashMap<String, CachedSchema>> = std::sync::RwLock::new(HashMap::new());
}

#[pg_extern(strict)]
fn cache_json_schemas(enums: JsonB, types: JsonB, puncs: JsonB) -> JsonB {
    let mut meta = SCHEMA_META.write().unwrap();
    let enums_value: Value = enums.0;
    let types_value: Value = types.0;
    let puncs_value: Value = puncs.0;

    let mut schemas_to_register = Vec::new();

    // Phase 1: Enums
    if let Some(enums_array) = enums_value.as_array() {
        for enum_row in enums_array {
            if let Some(schemas_raw) = enum_row.get("schemas") {
                if let Some(schemas_array) = schemas_raw.as_array() {
                    for schema_def in schemas_array {
                        if let Some(schema_id) = schema_def.get("$id").and_then(|v| v.as_str()) {
                            schemas_to_register.push((schema_id.to_string(), schema_def.clone(), SchemaType::Enum));
                        }
                    }
                }
            }
        }
    }

    // Phase 2: Types & Hierarchy
    let mut hierarchy_map: HashMap<String, HashSet<String>> = HashMap::new();
    if let Some(types_array) = types_value.as_array() {
        for type_row in types_array {
            if let Some(schemas_raw) = type_row.get("schemas") {
                if let Some(schemas_array) = schemas_raw.as_array() {
                    for schema_def in schemas_array {
                        if let Some(schema_id) = schema_def.get("$id").and_then(|v| v.as_str()) {
                            schemas_to_register.push((schema_id.to_string(), schema_def.clone(), SchemaType::Type));
                        }
                    }
                }
            }
            if let Some(type_name) = type_row.get("name").and_then(|v| v.as_str()) {
                if let Some(hierarchy_raw) = type_row.get("hierarchy") {
                    if let Some(hierarchy_array) = hierarchy_raw.as_array() {
                        for ancestor_val in hierarchy_array {
                            if let Some(ancestor_name) = ancestor_val.as_str() {
                                hierarchy_map.entry(ancestor_name.to_string()).or_default().insert(type_name.to_string());
                            }
                        }
                    }
                }
            }
        }
    }

    for (base_type, descendant_types) in hierarchy_map {
        let family_id = format!("{}.family", base_type);
        let values: Vec<String> = descendant_types.into_iter().collect();
        let family_schema = json!({ "$id": family_id, "type": "string", "enum": values });
        schemas_to_register.push((family_id, family_schema, SchemaType::Family));
    }

    // Phase 3: Puncs
    if let Some(puncs_array) = puncs_value.as_array() {
        for punc_row in puncs_array {
            if let Some(punc_obj) = punc_row.as_object() {
                if let Some(punc_name) = punc_obj.get("name").and_then(|v| v.as_str()) {
                    let is_public = punc_obj.get("public").and_then(|v| v.as_bool()).unwrap_or(false);
                    let punc_type = if is_public { SchemaType::PublicPunc } else { SchemaType::PrivatePunc };
                    if let Some(schemas_raw) = punc_obj.get("schemas") {
                        if let Some(schemas_array) = schemas_raw.as_array() {
                            for schema_def in schemas_array {
                                if let Some(schema_id) = schema_def.get("$id").and_then(|v| v.as_str()) {
                                    let req_id = format!("{}.request", punc_name);
                                    let resp_id = format!("{}.response", punc_name);
                                    let st = if schema_id == req_id || schema_id == resp_id { punc_type } else { SchemaType::Type };
                                    schemas_to_register.push((schema_id.to_string(), schema_def.clone(), st));
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    let mut all_errors = Vec::new();
    for (id, value, st) in schemas_to_register {
        // Meta-validation: Check 'type' enum if present
        if let Some(type_val) = value.get("type") {
            let types = match type_val {
                Value::String(s) => vec![s.as_str()],
                Value::Array(a) => a.iter().filter_map(|v| v.as_str()).collect(),
                _ => vec![],
            };
            let valid_primitives = ["string", "number", "integer", "boolean", "array", "object", "null"];
            for t in types {
                if !valid_primitives.contains(&t) {
                    all_errors.push(json!({ "code": "ENUM_VIOLATED", "message": format!("Invalid type: {}", t) }));
                }
            }
        }

        // Clone value for insertion since it might be consumed/moved if we were doing other things
        let value_for_registry = value.clone();
        
        // Validation: just ensure it is an object or boolean
        if value.is_object() || value.is_boolean() {
             REGISTRY.insert(id.clone(), value_for_registry);
             meta.insert(id, CachedSchema { t: st });
        } else {
             all_errors.push(json!({ "code": "INVALID_SCHEMA_TYPE", "message": format!("Schema {} must be an object or boolean", id) }));
        }
    }

    if !all_errors.is_empty() {
        return JsonB(json!({ "errors": all_errors }));
    }

    JsonB(json!({ "response": "success" }))
}

#[pg_extern(strict, parallel_safe)]
fn validate_json_schema(schema_id: &str, instance: JsonB) -> JsonB {
    let schema = match REGISTRY.get(schema_id) {
        Some(s) => s,
        None => return JsonB(json!({ 
            "errors": [{
                "code": "SCHEMA_NOT_FOUND",
                "message": format!("Schema '{}' not found", schema_id),
                "details": { "schema": schema_id }
            }]
        })),
    };

    let meta = SCHEMA_META.read().unwrap();
    let st = meta.get(schema_id).map(|m| m.t).unwrap_or(SchemaType::Type);
    
    let be_strict = match st {
        SchemaType::PublicPunc => true,
        _ => false,
    };
    
    let options = ValidationOptions {
        be_strict,
    };

    let mut validator = Validator::new(options, schema_id);
    match validator.validate(&schema, &instance.0) {
        Ok(_) => JsonB(json!({ "response": "success" })),
        Err(errors) => {
            let drop_errors: Vec<Value> = errors.into_iter().map(|e| json!({
                "code": e.code,
                "message": e.message,
                "details": {
                    "path": e.path,
                    "context": e.context,
                    "cause": e.cause,
                    "schema": e.schema_id
                }
            })).collect();
            
            if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open("/tmp/debug_jspg_errors.log") {
                 use std::io::Write;
                 let _ = writeln!(f, "VALIDATION FAILED for {}: {:?}", schema_id, drop_errors);
            }

            JsonB(json!({ "errors": drop_errors }))
        }
    }
}


#[pg_extern(strict, parallel_safe)]
fn json_schema_cached(schema_id: &str) -> bool {
    REGISTRY.get(schema_id).is_some()
}

#[pg_extern(strict)]
fn clear_json_schemas() -> JsonB {
    REGISTRY.reset();
    let mut meta = SCHEMA_META.write().unwrap();
    meta.clear();
    JsonB(json!({ "response": "success" }))
}

#[pg_extern(strict, parallel_safe)]
fn show_json_schemas() -> JsonB {
    let meta = SCHEMA_META.read().unwrap();
    let ids: Vec<String> = meta.keys().cloned().collect();
    JsonB(json!({ "response": ids }))
}

/// This module is required by `cargo pgrx test` invocations.
/// It must be visible at the root of your extension crate.
#[cfg(test)]
pub mod pg_test {
 pub fn setup(_options: Vec<&str>) {
  // perform one-off initialization when the pg_test framework starts
 }

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

#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
  use pgrx::pg_test;
 include!("suite.rs");
}