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

fn main() {
    let path = "tests/fixtures/lineage/snowflake_ground_truth.jsonl";
    let file = std::fs::File::open(path).unwrap();
    let reader = std::io::BufReader::new(file);

    let out_path = "scripts/lineage_review_data.jsonl";
    let mut out = std::io::BufWriter::new(std::fs::File::create(out_path).unwrap());

    let mut total = 0;
    let mut written = 0;
    let mut exact = 0;

    for line in reader.lines() {
        let line = line.unwrap();
        let rec: serde_json::Value = serde_json::from_str(&line).unwrap();
        total += 1;

        let sql = rec["sql"].as_str().unwrap();
        let name = rec["name"].as_str().unwrap_or("unknown");
        let query_type = rec
            .get("source_query_type")
            .and_then(|v| v.as_str())
            .unwrap_or("CTAS");

        let expected: BTreeMap<String, BTreeSet<String>> = rec["expected"]
            .as_object()
            .map(|map| {
                map.iter()
                    .map(|(k, v)| {
                        let vals: BTreeSet<String> = v
                            .as_array()
                            .map(|a| {
                                a.iter()
                                    .filter_map(|s| s.as_str().map(String::from))
                                    .collect()
                            })
                            .unwrap_or_default();
                        (k.clone(), vals)
                    })
                    .collect()
            })
            .unwrap_or_default();

        let schema: HashMap<String, serde_json::Value> =
            serde_json::from_value(rec["schema"].clone()).unwrap_or_default();

        let result = match column_lineage(sql, SqlDialect::Snowflake, &schema, None) {
            Ok(r) => r,
            Err(e) => {
                let record = serde_json::json!({
                    "name": name,
                    "sql": sql,
                    "query_type": query_type,
                    "expected": &expected,
                    "actual": {},
                    "category": "parse_error",
                    "error": e,
                });
                writeln!(out, "{}", serde_json::to_string(&record).unwrap()).unwrap();
                written += 1;
                continue;
            }
        };

        let actual: BTreeMap<String, BTreeSet<String>> = result
            .column_dict
            .iter()
            .map(|(k, v)| {
                let vals: BTreeSet<String> = v.iter().cloned().collect();
                (k.clone(), vals)
            })
            .collect();

        if actual
            == expected
                .iter()
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect::<BTreeMap<_, _>>()
        {
            exact += 1;
            continue; // Skip exact matches
        }

        // Classify
        let has_extra_cols = actual.keys().any(|k| !expected.contains_key(k));
        let has_missing_cols = expected.keys().any(|k| !actual.contains_key(k));
        let has_wrong_sources = expected.keys().any(|k| {
            actual.get(k).map_or(false, |act| {
                let exp: &BTreeSet<String> = expected.get(k).unwrap();
                act != exp
            })
        });

        let same_cols = !has_extra_cols && !has_missing_cols;
        let is_superset = same_cols
            && expected.iter().all(|(col, exp_srcs)| {
                actual
                    .get(col)
                    .map_or(false, |act_srcs| exp_srcs.is_subset(act_srcs))
            });

        let category = if is_superset {
            "superset_only"
        } else if has_extra_cols && !has_missing_cols && !has_wrong_sources {
            "extra_cols_only"
        } else if !has_extra_cols && has_missing_cols && !has_wrong_sources {
            "missing_cols_only"
        } else if same_cols {
            "wrong_sources"
        } else {
            "mixed"
        };

        // Convert to serializable format
        let actual_ser: BTreeMap<String, Vec<String>> = actual
            .iter()
            .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
            .collect();
        let expected_ser: BTreeMap<String, Vec<String>> = expected
            .iter()
            .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
            .collect();

        let record = serde_json::json!({
            "name": name,
            "sql": sql,
            "query_type": query_type,
            "expected": expected_ser,
            "actual": actual_ser,
            "category": category,
        });
        writeln!(out, "{}", serde_json::to_string(&record).unwrap()).unwrap();
        written += 1;
    }

    eprintln!(
        "Done: {}/{} total, {} exact, {} failures written to {}",
        total, total, exact, written, out_path
    );
}
