use std::collections::BTreeMap;

use serde_json::Value;

use crate::{JsonPathRoot, MappingRule, MappingSet};

pub fn infer_mapping_from_samples(
    source_sample: &Value,
    source_root: JsonPathRoot,
    target_sample: &Value,
    target_root: JsonPathRoot,
) -> MappingSet {
    let source_paths = collect_leaf_paths(source_sample, source_root);
    let target_paths = collect_leaf_paths(target_sample, target_root);
    let mut rules = Vec::new();

    for (leaf_name, source_path) in source_paths {
        let Some(target_path) = target_paths.get(&leaf_name) else {
            continue;
        };

        rules.push(MappingRule {
            source: source_path,
            target: target_path.clone(),
            required: false,
            default_value: None,
            transform: None,
            condition: None,
            notes: None,
        });
    }

    MappingSet { rules }
}

fn collect_leaf_paths(sample: &Value, root: JsonPathRoot) -> BTreeMap<String, String> {
    let mut collected = BTreeMap::new();
    let mut segments = Vec::new();

    collect_leaf_paths_recursive(sample, &mut segments, &mut collected);

    collected
        .into_iter()
        .map(|(leaf, path_segments)| {
            let mut path = format!("$.{}", root.as_str());

            for segment in path_segments {
                path.push('.');
                path.push_str(&segment);
            }

            (leaf, path)
        })
        .collect()
}

fn collect_leaf_paths_recursive(
    sample: &Value,
    segments: &mut Vec<String>,
    collected: &mut BTreeMap<String, Vec<String>>,
) {
    match sample {
        Value::Object(object) => {
            for (field, value) in object {
                segments.push(field.clone());
                collect_leaf_paths_recursive(value, segments, collected);
                segments.pop();
            }
        }
        Value::Array(values) => {
            if let Some(first) = values.first() {
                collect_leaf_paths_recursive(first, segments, collected);
            }
        }
        Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {
            if let Some(leaf) = segments.last() {
                collected
                    .entry(leaf.clone())
                    .or_insert_with(|| segments.clone());
            }
        }
    }
}

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

    use crate::{JsonPathRoot, infer_mapping_from_samples};

    #[test]
    fn infers_rules_from_matching_leaf_names() {
        let mapping = infer_mapping_from_samples(
            &json!({
                "lead": {
                    "email": "user@example.com",
                    "name": "Ada"
                }
            }),
            JsonPathRoot::Mcp,
            &json!({
                "contact": {
                    "email": "user@example.com"
                },
                "name": "Ada"
            }),
            JsonPathRoot::RequestBody,
        );

        assert_eq!(mapping.rules.len(), 2);
        assert_eq!(mapping.rules[0].source, "$.mcp.lead.email");
        assert_eq!(mapping.rules[0].target, "$.request.body.contact.email");
        assert_eq!(mapping.rules[1].source, "$.mcp.lead.name");
        assert_eq!(mapping.rules[1].target, "$.request.body.name");
    }
}
