use serde_json::Value;
use std::collections::HashMap;
use std::sync::RwLock;
use lazy_static::lazy_static;

lazy_static! {
    pub static ref REGISTRY: Registry = Registry::new();
}

pub struct Registry {
    schemas: RwLock<HashMap<String, Value>>,
}

impl Registry {
    pub fn new() -> Self {
        Self {
            schemas: RwLock::new(HashMap::new()),
        }
    }

    pub fn reset(&self) {
        let mut schemas = self.schemas.write().unwrap();
        schemas.clear();
    }

    pub fn insert(&self, id: String, schema: Value) {
        let mut schemas = self.schemas.write().unwrap();
        
        // Index the schema and its sub-resources (IDs and anchors)
        self.index_schema(&schema, &mut schemas, Some(&id));
        
        // Ensure the root ID is inserted (index_schema handles it, but let's be explicit)
        schemas.insert(id, schema);
    }

    fn index_schema(&self, schema: &Value, registry: &mut HashMap<String, Value>, current_scope: Option<&str>) {
        if let Value::Object(map) = schema {
            // Only strictly index $id for scope resolution
            let mut my_scope = current_scope.map(|s| s.to_string());

            if let Some(Value::String(id)) = map.get("$id") {
                if id.contains("://") {
                    my_scope = Some(id.clone());
                } else if let Some(scope) = current_scope {
                    if let Some(pos) = scope.rfind('/') {
                        my_scope = Some(format!("{}{}", &scope[..pos + 1], id));
                    } else {
                        my_scope = Some(id.clone());
                    }
                } else {
                    my_scope = Some(id.clone());
                }

                if let Some(final_id) = &my_scope {
                    registry.insert(final_id.clone(), schema.clone());
                }
            }

            // Minimal recursion only for definitions where sub-IDs often live
            // This is a tradeoff: we don't index EVERYWHERE, but we catch the 90% common case of
            // bundled definitions without full tree traversal.
            if let Some(Value::Object(defs)) = map.get("$defs").or_else(|| map.get("definitions")) {
                for (_, def_schema) in defs {
                    self.index_schema(def_schema, registry, my_scope.as_deref());
                }
            }
        }
    }

    pub fn get(&self, id: &str) -> Option<Value> {
        let schemas = self.schemas.read().unwrap();
        schemas.get(id).cloned()
    }

    pub fn resolve(&self, ref_str: &str, current_id: Option<&str>) -> Option<(Value, String)> {
        // 1. Try full lookup (Absolute or explicit ID)
        if let Some(s) = self.get(ref_str) {
            return Some((s, ref_str.to_string()));
        }

        // 2. Try Relative lookup against current scope
        if let Some(curr) = current_id {
            if let Some(pos) = curr.rfind('/') {
                let joined = format!("{}{}", &curr[..pos + 1], ref_str);
                if let Some(s) = self.get(&joined) {
                    return Some((s, joined));
                }
            }
        }

        // 3. Pointer Resolution
        // Split into Base URI + Fragment
        let (base, fragment) = match ref_str.split_once('#') {
            Some((b, f)) => (b, Some(f)),
            None => (ref_str, None),
        };

        // If base is empty, we stay in current schema.
        // If base is present, we resolve it first.
        let (root_schema, scope) = if base.is_empty() {
            if let Some(curr) = current_id {
                // If we are looking up internally, we rely on the caller having passed the correct current ID
                // But typically internal refs are just fragments.
                if let Some(s) = self.get(curr) {
                    (s, curr.to_string())
                } else {
                    return None;
                }
            } else {
                return None;
            }
        } else {
            // Resolve external base
            if let Some(s) = self.get(base) {
                (s, base.to_string())
            } else if let Some(curr) = current_id {
                 // Try relative base
                 if let Some(pos) = curr.rfind('/') {
                    let joined = format!("{}{}", &curr[..pos + 1], base);
                    if let Some(s) = self.get(&joined) {
                        (s, joined)
                    } else {
                        return None;
                    }
                 } else {
                     return None;
                 }
            } else {
                return None;
            }
        };

        if let Some(frag_raw) = fragment {
            if frag_raw.is_empty() {
                return Some((root_schema, scope));
            }
            // Decode fragment (it is URI encoded)
            let frag_cow = percent_encoding::percent_decode_str(frag_raw).decode_utf8().unwrap_or(std::borrow::Cow::Borrowed(frag_raw));
            let frag = frag_cow.as_ref();

            if frag.starts_with('/') {
                if let Some(sub) = root_schema.pointer(frag) {
                    return Some((sub.clone(), scope));
                }
            } else {
                // It is an anchor. We scan for it at runtime to avoid complex indexing at insertion.
                if let Some(sub) = self.find_anchor(&root_schema, frag) {
                    return Some((sub, scope));
                }
            }
            None
        } else {
            Some((root_schema, scope))
        }
    }

    fn find_anchor(&self, schema: &Value, anchor: &str) -> Option<Value> {
        match schema {
            Value::Object(map) => {
                // Check if this schema itself has the anchor
                if let Some(Value::String(a)) = map.get("$anchor") {
                    if a == anchor {
                        return Some(schema.clone());
                    }
                }
                
                // Recurse into $defs / definitions (Map of Schemas)
                if let Some(Value::Object(defs)) = map.get("$defs").or_else(|| map.get("definitions")) {
                    for val in defs.values() {
                        if let Some(found) = self.find_anchor(val, anchor) { return Some(found); }
                    }
                }

                // Recurse into properties / patternProperties / dependentSchemas (Map of Schemas)
                for key in ["properties", "patternProperties", "dependentSchemas"] {
                    if let Some(Value::Object(props)) = map.get(key) {
                        for val in props.values() {
                            if let Some(found) = self.find_anchor(val, anchor) { return Some(found); }
                        }
                    }
                }

                // Recurse into arrays of schemas
                for key in ["allOf", "anyOf", "oneOf", "prefixItems"] {
                    if let Some(Value::Array(arr)) = map.get(key) {
                        for item in arr {
                            if let Some(found) = self.find_anchor(item, anchor) { return Some(found); }
                        }
                    }
                }

                // Recurse into single sub-schemas
                for key in ["items", "contains", "additionalProperties", "unevaluatedProperties", "not", "if", "then", "else"] {
                    if let Some(val) = map.get(key) {
                         if val.is_object() || val.is_boolean() {
                             if let Some(found) = self.find_anchor(val, anchor) { return Some(found); }
                         }
                    }
                }
                None
            }
            Value::Array(arr) => {
                // Should not happen for a schema object, but if we are passed an array of schemas?
                // Standard schema is object or bool. 
                // But let's be safe.
                for item in arr {
                    if let Some(found) = self.find_anchor(item, anchor) {
                        return Some(found);
                    }
                }
                None
            }
            _ => None,
        }
    }
}