use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use std::collections::{BTreeSet, HashMap};

fn run_case(label: &str, path: &str) {
    let data = std::fs::read_to_string(path).unwrap();
    let case: serde_json::Value = serde_json::from_str(&data).unwrap();

    let sql = case["sql"].as_str().unwrap();
    let schema: HashMap<String, serde_json::Value> = case["schema"]
        .as_object()
        .unwrap()
        .iter()
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();
    let expected: HashMap<String, Vec<String>> =
        serde_json::from_value(case["expected"].clone()).unwrap();

    let result = match column_lineage(sql, SqlDialect::Snowflake, &schema, None) {
        Ok(r) => r,
        Err(e) => {
            println!("=== {label}: ERROR ===\n{e}\n");
            return;
        }
    };

    println!("=== {label} ===");

    // Compare each column
    let mut all_cols: BTreeSet<String> = BTreeSet::new();
    all_cols.extend(expected.keys().cloned());
    all_cols.extend(result.column_dict.keys().cloned());

    let mut wrong = 0;
    for col in &all_cols {
        let exp: BTreeSet<String> = expected
            .get(col)
            .map(|v| v.iter().cloned().collect())
            .unwrap_or_default();
        let act: BTreeSet<String> = result
            .column_dict
            .get(col)
            .map(|v| {
                v.iter()
                    .map(|s| {
                        let parts: Vec<&str> = s.split('.').map(|p| p.trim_matches('"')).collect();
                        if parts.len() >= 3 {
                            let t = parts[parts.len() - 2];
                            let c = parts[parts.len() - 1];
                            format!("\"{t}\".\"{c}\"")
                        } else {
                            s.clone()
                        }
                    })
                    .collect()
            })
            .unwrap_or_default();

        if exp == act {
            continue; // correct
        }

        wrong += 1;
        let extras: BTreeSet<_> = act.difference(&exp).cloned().collect();
        let missing: BTreeSet<_> = exp.difference(&act).cloned().collect();

        print!("  {col}: ");
        if !extras.is_empty() && missing.is_empty() {
            println!("SUPERSET +{:?}", extras);
        } else if extras.is_empty() && !missing.is_empty() {
            println!("SUBSET -{:?}", missing);
        } else if exp.is_empty() {
            println!("EXTRA COL (not expected), act={:?}", act);
        } else if act.is_empty() {
            println!("MISSING COL (no actual), exp={:?}", exp);
        } else {
            println!(
                "MISMATCH\n    extra: {:?}\n    missing: {:?}",
                extras, missing
            );
        }
    }
    if wrong == 0 {
        println!("  ALL CORRECT!");
    }
    println!();
}

fn main() {
    run_case("STORES.TENANT leak", "/tmp/test_case_1.json");
    run_case("GLOBAL_AUDIENCE leak", "/tmp/test_case_2.json");
    run_case("CASE WHEN condition bleed", "/tmp/test_case_3.json");
    run_case("Window PARTITION BY bleed", "/tmp/test_case_4.json");
}
