//! Local JSON Schema reference resolution and `allOf` merging for argument
//! synthesis.
//!
//! Why this exists: real tool catalogs factor shared shapes into `$defs` (or
//! the older `definitions`) and compose them with `allOf`, so synthesizing
//! arguments from the raw property schema alone would miss most constraints.
//! This module flattens those two mechanisms into one effective schema the
//! synthesis rules in [`super::synthesis`] can read directly. Only local
//! (`#/...`) references resolve; a generator that must work offline cannot
//! chase remote refs.

use serde_json::Value;

/// How many `$ref` hops we chase before giving up. A required
/// self-referential schema would otherwise recurse forever; eight levels is
/// deeper than any real catalog while keeping the failure mode a YAML
/// comment, not a hang.
pub(crate) const MAX_REF_DEPTH: usize = 8;

/// Flatten `$ref` chains and `allOf` branches into one effective schema.
///
/// `root` is the document `#/...` pointers resolve against (the tool's whole
/// `inputSchema`). `depth` counts `$ref` hops so callers can thread the
/// budget through nested recursion; the returned depth is the consumed
/// total. Errors are human-readable reasons (unresolvable ref, depth cap,
/// contradictory `allOf` types) that the generators surface as comments.
pub(crate) fn resolve(
    schema: &Value,
    root: &Value,
    depth: usize,
    notes: &mut Vec<String>,
) -> Result<(Value, usize), String> {
    let mut current = schema.clone();
    let mut depth = depth;
    while let Some(reference) = current
        .get("$ref")
        .and_then(Value::as_str)
        .map(str::to_string)
    {
        if depth >= MAX_REF_DEPTH {
            return Err(format!(
                "reference depth cap ({MAX_REF_DEPTH}) reached at `{reference}`"
            ));
        }
        depth += 1;
        let mut siblings = current.as_object().cloned().unwrap_or_default();
        siblings.remove("$ref");
        let Some(pointer) = reference.strip_prefix('#') else {
            notes.push(format!(
                "remote $ref `{reference}` is not resolved; treating it as unconstrained"
            ));
            current = Value::Object(siblings);
            continue;
        };
        let Some(target) = root.pointer(pointer) else {
            return Err(format!("local $ref `{reference}` does not resolve"));
        };
        current = if siblings.is_empty() {
            target.clone()
        } else {
            merge(target, &Value::Object(siblings))?
        };
    }
    fold_all_of(current, root, depth, notes)
}

/// Merge every `allOf` branch into the parent schema. Branches resolve
/// recursively first so a `$ref` inside an `allOf` still lands in the fold.
fn fold_all_of(
    current: Value,
    root: &Value,
    depth: usize,
    notes: &mut Vec<String>,
) -> Result<(Value, usize), String> {
    let Some(branches) = current.get("allOf").and_then(Value::as_array).cloned() else {
        return Ok((current, depth));
    };
    let mut base = current.as_object().cloned().unwrap_or_default();
    base.remove("allOf");
    let mut folded = Value::Object(base);
    let mut max_depth = depth;
    for branch in &branches {
        let (resolved, branch_depth) = resolve(branch, root, depth, notes)?;
        max_depth = max_depth.max(branch_depth);
        folded = merge(&folded, &resolved)?;
    }
    Ok((folded, max_depth))
}

/// Combine two schema objects under `allOf` semantics. Properties and
/// required lists union, numeric and length bounds tighten, conflicting
/// `type` keywords are an error, and every other keyword keeps the first
/// branch's value (boring and predictable beats a speculative full merge).
pub(crate) fn merge(a: &Value, b: &Value) -> Result<Value, String> {
    let mut out = a.as_object().cloned().unwrap_or_default();
    let Some(extra) = b.as_object() else {
        return Ok(Value::Object(out));
    };
    for (key, value) in extra {
        let merged = match out.get(key) {
            None => value.clone(),
            Some(existing) if existing == value => existing.clone(),
            Some(existing) => merge_keyword(key, existing, value)?,
        };
        out.insert(key.clone(), merged);
    }
    Ok(Value::Object(out))
}

fn merge_keyword(key: &str, ours: &Value, theirs: &Value) -> Result<Value, String> {
    match key {
        "type" => merge_types(ours, theirs),
        "properties" => Ok(merge_properties(ours, theirs)),
        "required" => Ok(merge_required(ours, theirs)),
        "minimum" | "exclusiveMinimum" | "minLength" | "minItems" => {
            Ok(tighter_bound(ours, theirs, true))
        }
        "maximum" | "exclusiveMaximum" | "maxLength" | "maxItems" => {
            Ok(tighter_bound(ours, theirs, false))
        }
        _ => Ok(ours.clone()),
    }
}

fn merge_types(ours: &Value, theirs: &Value) -> Result<Value, String> {
    let a = type_list(ours);
    let b = type_list(theirs);
    if a.is_empty() {
        return Ok(theirs.clone());
    }
    if b.is_empty() {
        return Ok(ours.clone());
    }
    let mut intersection: Vec<String> = a.iter().filter(|t| b.contains(t)).cloned().collect();
    if intersection.is_empty() {
        // `integer` narrows `number` rather than conflicting with it.
        let narrows = (a.contains(&"number".to_string()) && b.contains(&"integer".to_string()))
            || (b.contains(&"number".to_string()) && a.contains(&"integer".to_string()));
        if !narrows {
            return Err(format!(
                "allOf branches declare conflicting types: `{}` vs `{}`",
                a.join("|"),
                b.join("|")
            ));
        }
        intersection.push("integer".to_string());
    }
    if intersection.len() == 1 {
        Ok(Value::String(intersection.remove(0)))
    } else {
        Ok(Value::Array(
            intersection.into_iter().map(Value::String).collect(),
        ))
    }
}

fn type_list(value: &Value) -> Vec<String> {
    match value {
        Value::String(s) => vec![s.clone()],
        Value::Array(arr) => arr
            .iter()
            .filter_map(Value::as_str)
            .map(str::to_string)
            .collect(),
        _ => Vec::new(),
    }
}

fn merge_properties(ours: &Value, theirs: &Value) -> Value {
    let mut out = ours.as_object().cloned().unwrap_or_default();
    let Some(extra) = theirs.as_object() else {
        return Value::Object(out);
    };
    for (name, schema) in extra {
        match out.get(name) {
            None => {
                out.insert(name.clone(), schema.clone());
            }
            Some(existing) if existing == schema => {}
            Some(existing) => {
                // Both branches constrain the same property; keep both via a
                // nested allOf so synthesis re-merges (and conflict-checks)
                // when the property is actually synthesized.
                let combined = serde_json::json!({ "allOf": [existing, schema] });
                out.insert(name.clone(), combined);
            }
        }
    }
    Value::Object(out)
}

fn merge_required(ours: &Value, theirs: &Value) -> Value {
    let mut names: Vec<Value> = ours.as_array().cloned().unwrap_or_default();
    for entry in theirs.as_array().into_iter().flatten() {
        if !names.contains(entry) {
            names.push(entry.clone());
        }
    }
    Value::Array(names)
}

fn tighter_bound(ours: &Value, theirs: &Value, prefer_max: bool) -> Value {
    let (Some(a), Some(b)) = (ours.as_f64(), theirs.as_f64()) else {
        // Boolean-style draft-4 exclusives (or junk) fall back to first-wins.
        return ours.clone();
    };
    let take_theirs = if prefer_max { b > a } else { b < a };
    if take_theirs {
        theirs.clone()
    } else {
        ours.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn resolve_ok(schema: &Value, root: &Value) -> Value {
        let mut notes = Vec::new();
        resolve(schema, root, 0, &mut notes).expect("resolves").0
    }

    #[test]
    fn local_ref_into_defs_resolves_to_target() {
        let root = json!({
            "$defs": { "city": { "type": "string", "minLength": 2 } },
            "properties": { "city": { "$ref": "#/$defs/city" } }
        });
        let resolved = resolve_ok(&root["properties"]["city"], &root);
        assert_eq!(resolved["type"], "string");
        assert_eq!(resolved["minLength"], 2);
    }

    #[test]
    fn local_ref_into_definitions_resolves_too() {
        let root = json!({
            "definitions": { "n": { "type": "integer" } }
        });
        let resolved = resolve_ok(&json!({ "$ref": "#/definitions/n" }), &root);
        assert_eq!(resolved["type"], "integer");
    }

    #[test]
    fn unresolvable_local_ref_is_an_error() {
        let root = json!({});
        let mut notes = Vec::new();
        let err = resolve(&json!({ "$ref": "#/$defs/missing" }), &root, 0, &mut notes)
            .expect_err("missing pointer fails");
        assert!(err.contains("does not resolve"), "{err}");
    }

    #[test]
    fn self_referential_chain_stops_at_the_depth_cap() {
        let root = json!({
            "$defs": { "loop": { "$ref": "#/$defs/loop" } }
        });
        let mut notes = Vec::new();
        let err = resolve(&json!({ "$ref": "#/$defs/loop" }), &root, 0, &mut notes)
            .expect_err("infinite chain hits the cap");
        assert!(err.contains("depth cap"), "{err}");
    }

    #[test]
    fn remote_ref_is_noted_and_treated_as_unconstrained() {
        let root = json!({});
        let mut notes = Vec::new();
        let (resolved, _) = resolve(
            &json!({ "$ref": "https://example.com/s.json", "description": "d" }),
            &root,
            0,
            &mut notes,
        )
        .expect("remote ref degrades gracefully");
        assert!(resolved.get("$ref").is_none());
        assert_eq!(resolved["description"], "d");
        assert_eq!(notes.len(), 1);
        assert!(notes[0].contains("remote $ref"));
    }

    #[test]
    fn all_of_unions_properties_and_required() {
        let schema = json!({
            "allOf": [
                { "type": "object", "properties": { "a": { "type": "string" } }, "required": ["a"] },
                { "type": "object", "properties": { "b": { "type": "integer" } }, "required": ["b"] }
            ]
        });
        let resolved = resolve_ok(&schema, &schema);
        assert_eq!(resolved["properties"]["a"]["type"], "string");
        assert_eq!(resolved["properties"]["b"]["type"], "integer");
        assert_eq!(resolved["required"], json!(["a", "b"]));
    }

    #[test]
    fn all_of_type_conflict_is_an_error() {
        let schema = json!({
            "allOf": [{ "type": "string" }, { "type": "integer" }]
        });
        let mut notes = Vec::new();
        let err =
            resolve(&schema, &schema, 0, &mut notes).expect_err("string vs integer conflicts");
        assert!(err.contains("conflicting types"), "{err}");
    }

    #[test]
    fn all_of_integer_narrows_number_instead_of_conflicting() {
        let schema = json!({
            "allOf": [{ "type": "number" }, { "type": "integer" }]
        });
        let resolved = resolve_ok(&schema, &schema);
        assert_eq!(resolved["type"], "integer");
    }

    #[test]
    fn all_of_tightens_bounds_to_the_narrowest_range() {
        let schema = json!({
            "allOf": [
                { "type": "integer", "minimum": 1, "maximum": 100 },
                { "minimum": 10, "maximum": 50 }
            ]
        });
        let resolved = resolve_ok(&schema, &schema);
        assert_eq!(resolved["minimum"], 10);
        assert_eq!(resolved["maximum"], 50);
    }

    #[test]
    fn ref_sibling_keywords_merge_onto_the_target() {
        let root = json!({
            "$defs": { "s": { "type": "string" } }
        });
        let resolved = resolve_ok(&json!({ "$ref": "#/$defs/s", "minLength": 3 }), &root);
        assert_eq!(resolved["type"], "string");
        assert_eq!(resolved["minLength"], 3);
    }
}
