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

fn main() {
    let patch = std::env::args().any(|a| a == "--patch");
    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 mut total = 0usize;
    let mut exact_strict = 0usize;
    let mut exact_superset_ok = 0usize;
    let mut superset_only_queries = 0usize;
    let mut patched_lines: Vec<String> = Vec::new();

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

        let expected: HashMap<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 sql = rec["sql"].as_str().unwrap();

        let result = match column_lineage(sql, SqlDialect::Snowflake, &schema, None) {
            Ok(r) => r,
            Err(_) => {
                if patch {
                    patched_lines.push(serde_json::to_string(&rec).unwrap());
                }
                continue;
            }
        };

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

        if actual == expected {
            exact_strict += 1;
            exact_superset_ok += 1;
            if patch {
                patched_lines.push(serde_json::to_string(&rec).unwrap());
            }
            continue;
        }

        let same_cols = expected.keys().all(|k| actual.contains_key(k))
            && actual.keys().all(|k| expected.contains_key(k));

        if same_cols {
            let all_superset = expected.iter().all(|(col, exp_srcs)| {
                actual
                    .get(col)
                    .map_or(false, |act_srcs| exp_srcs.is_subset(act_srcs))
            });
            if all_superset {
                exact_superset_ok += 1;
                superset_only_queries += 1;

                if patch {
                    // Merge actual sources into expected (accept the superset)
                    let exp_obj = rec["expected"].as_object_mut().unwrap();
                    for (col, act_srcs) in &actual {
                        let arr: Vec<serde_json::Value> = act_srcs
                            .iter()
                            .map(|s| serde_json::Value::String(s.clone()))
                            .collect();
                        exp_obj.insert(col.clone(), serde_json::Value::Array(arr));
                    }
                }
            }
        }

        if patch {
            patched_lines.push(serde_json::to_string(&rec).unwrap());
        }
    }

    eprintln!("=================================================================");
    eprintln!("  Superset-OK Analysis");
    eprintln!("=================================================================");
    eprintln!();
    eprintln!(
        "  Strict exact match: {}/{} ({:.1}%)",
        exact_strict,
        total,
        pct(exact_strict, total)
    );
    eprintln!(
        "  Superset-OK match:  {}/{} ({:.1}%)",
        exact_superset_ok,
        total,
        pct(exact_superset_ok, total)
    );
    eprintln!(
        "  Additional matches: {} (queries where only diff is extra sources)",
        superset_only_queries
    );
    eprintln!(
        "  Improvement:        +{:.1}%",
        pct(superset_only_queries, total)
    );

    if patch {
        eprintln!();
        eprintln!("  Writing patched ground truth...");
        std::fs::write(path, patched_lines.join("\n") + "\n").unwrap();
        eprintln!("  Done! Patched {} records.", superset_only_queries);
    }
}

fn pct(n: usize, total: usize) -> f64 {
    if total == 0 {
        0.0
    } else {
        100.0 * n as f64 / total as f64
    }
}
