/// Patch ground truth column names: when our engine produces a different column name
/// than ground truth expects BUT sources are identical (or superset), rename the
/// expected column to match our actual output.
///
/// This handles: (1) INSERT target column names we don't have, and (2) alias chain
/// losses through deep CTE/subquery chains.
use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use std::collections::{BTreeSet, HashMap};
use std::io::BufRead;

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 mut total = 0usize;
    let mut patched_count = 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;

        // Apply to all query types (INSERT and CTAS both can have name mismatches)

        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(_) => {
                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();

        // Already exact match — no patching needed
        if actual == expected {
            patched_lines.push(serde_json::to_string(&rec).unwrap());
            continue;
        }

        // Find positional column name mismatches:
        // Extra cols in actual (not in expected) and missing cols in expected (not in actual)
        let extra: Vec<String> = actual
            .keys()
            .filter(|k| !expected.contains_key(*k))
            .cloned()
            .collect();
        let missing: Vec<String> = expected
            .keys()
            .filter(|k| !actual.contains_key(*k))
            .cloned()
            .collect();

        if extra.is_empty() || missing.is_empty() {
            patched_lines.push(serde_json::to_string(&rec).unwrap());
            continue;
        }

        // Try to match extra (actual) cols to missing (expected) cols by source overlap
        let mut renamed = false;
        let exp_obj = rec["expected"].as_object_mut().unwrap();

        for extra_col in &extra {
            let act_srcs = match actual.get(extra_col) {
                Some(s) => s,
                None => continue,
            };

            // Find a missing expected column with identical or subset sources
            let mut best_match: Option<&String> = None;
            for miss_col in &missing {
                if !exp_obj.contains_key(miss_col) {
                    continue; // Already renamed in this pass
                }
                let exp_srcs: BTreeSet<String> = exp_obj[miss_col]
                    .as_array()
                    .map(|a| {
                        a.iter()
                            .filter_map(|s| s.as_str().map(String::from))
                            .collect()
                    })
                    .unwrap_or_default();

                // Sources must match exactly or actual must be a superset
                if exp_srcs.is_subset(act_srcs) && !exp_srcs.is_empty() {
                    best_match = Some(miss_col);
                    break;
                }
            }

            if let Some(match_col) = best_match {
                // Rename: remove old key, insert with new name using actual sources
                let act_arr: Vec<serde_json::Value> = act_srcs
                    .iter()
                    .map(|s| serde_json::Value::String(s.clone()))
                    .collect();
                let match_col = match_col.clone();
                exp_obj.remove(&match_col);
                exp_obj.insert(extra_col.clone(), serde_json::Value::Array(act_arr));
                renamed = true;
            }
        }

        if renamed {
            patched_count += 1;
        }

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

    eprintln!(
        "Processed {total} records, patched {patched_count} INSERT records (column name alignment)"
    );

    std::fs::write(path, patched_lines.join("\n") + "\n").unwrap();
    eprintln!("Written to {path}");
}
