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 = 0;
    let mut exact = 0;

    // Root cause counters
    let mut rc_unnamed_col = 0; // _col_N naming issue
    let mut rc_star_expand = 0; // SELECT * expansion issue
    let mut rc_union_branch = 0; // UNION ALL branch tracing
    let mut rc_window_extra = 0; // Window function extra columns
    let mut rc_col_alias = 0; // Column alias not matching expected
    let mut rc_missing_src = 0; // Missing source columns
    let mut rc_extra_src = 0; // Extra source columns
    let mut rc_empty = 0; // Empty result
    let mut rc_other = 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 expected_map: HashMap<String, Vec<String>> =
            serde_json::from_value(rec["expected"].clone()).unwrap();
        let schema_val = rec
            .get("schema")
            .cloned()
            .unwrap_or(serde_json::Value::Object(Default::default()));
        let schema: HashMap<String, serde_json::Value> =
            serde_json::from_value(schema_val).unwrap();

        let result = match column_lineage(sql, SqlDialect::Snowflake, &schema, None) {
            Ok(r) => r,
            Err(_) => {
                rc_other += 1;
                continue;
            }
        };

        // Normalize both
        let mut actual_norm: HashMap<String, BTreeSet<String>> = HashMap::new();
        for (k, v) in &result.column_dict {
            let key = k.to_uppercase();
            let sources: BTreeSet<String> = v
                .iter()
                .map(|s| s.replace('"', "").to_uppercase())
                .collect();
            actual_norm.insert(key, sources);
        }

        let mut expected_norm: HashMap<String, BTreeSet<String>> = HashMap::new();
        for (k, v) in &expected_map {
            let key = k.to_uppercase();
            let sources: BTreeSet<String> = v
                .iter()
                .map(|s| s.replace('"', "").to_uppercase())
                .collect();
            expected_norm.insert(key, sources);
        }

        if actual_norm == expected_norm {
            exact += 1;
            continue;
        }

        // Categorize the failure
        let actual_keys: BTreeSet<&String> = actual_norm.keys().collect();
        let expected_keys: BTreeSet<&String> = expected_norm.keys().collect();

        let extra_keys: Vec<&&String> = actual_keys.difference(&expected_keys).collect();
        let missing_keys: Vec<&&String> = expected_keys.difference(&actual_keys).collect();

        let sql_upper = sql.to_uppercase();

        // Check for _col_N naming
        let has_unnamed = extra_keys.iter().any(|k| k.starts_with("_COL_"));

        // Check for SELECT *
        let has_star = sql_upper.contains("SELECT *") || sql_upper.contains("SELECT  *");

        // Check for UNION ALL
        let has_union = sql_upper.contains("UNION ALL") || sql_upper.contains("UNION\n");

        // Check for window functions in extra keys
        let window_keywords = ["ROW_NUMBER", "RANK", "DENSE_RANK", "NTILE", "LAG", "LEAD"];
        let has_window_extra = extra_keys.iter().any(|k| {
            // Check if this extra col seems window-related
            if let Some(srcs) = actual_norm.get(**k) {
                // Window functions typically have PARTITION BY sources
                srcs.len() > 1
            } else {
                false
            }
        }) && window_keywords.iter().any(|w| sql_upper.contains(w));

        // Categorize
        if actual_norm.is_empty() {
            rc_empty += 1;
        } else if has_unnamed && !missing_keys.is_empty() {
            rc_unnamed_col += 1;
        } else if has_star && (!missing_keys.is_empty() || !extra_keys.is_empty()) {
            rc_star_expand += 1;
        } else if has_union && !missing_keys.is_empty() {
            rc_union_branch += 1;
        } else if has_window_extra && extra_keys.len() > 0 && missing_keys.is_empty() {
            rc_window_extra += 1;
        } else if !extra_keys.is_empty() && !missing_keys.is_empty() {
            // Has both extra and missing keys - likely alias mismatch
            rc_col_alias += 1;
        } else if extra_keys.is_empty() && missing_keys.is_empty() {
            // Same keys but different sources
            let mut all_missing = false;
            let mut all_extra = false;
            let mut both = false;
            for k in actual_keys.intersection(&expected_keys) {
                let a = actual_norm.get(*k).unwrap();
                let e = expected_norm.get(*k).unwrap();
                if a != e {
                    let missing: Vec<_> = e.difference(a).collect();
                    let extra: Vec<_> = a.difference(e).collect();
                    if !missing.is_empty() && extra.is_empty() {
                        all_missing = true;
                    } else if missing.is_empty() && !extra.is_empty() {
                        all_extra = true;
                    } else {
                        both = true;
                    }
                }
            }
            if all_missing && !all_extra && !both {
                rc_missing_src += 1;
            } else if all_extra && !all_missing && !both {
                rc_extra_src += 1;
            } else {
                rc_other += 1;
            }
        } else if extra_keys.is_empty() && !missing_keys.is_empty() {
            rc_missing_src += 1;
        } else if !extra_keys.is_empty() && missing_keys.is_empty() {
            rc_extra_src += 1;
        } else {
            rc_other += 1;
        }
    }

    eprintln!("=== ROOT CAUSE ANALYSIS ({total} queries) ===");
    eprintln!(
        "  Exact match:      {exact:6} ({:.1}%)",
        exact as f64 * 100.0 / total as f64
    );
    eprintln!("  ---");
    eprintln!("  Empty result:     {rc_empty:6}");
    eprintln!("  Unnamed _col_N:   {rc_unnamed_col:6}");
    eprintln!("  SELECT * expand:  {rc_star_expand:6}");
    eprintln!("  UNION branch:     {rc_union_branch:6}");
    eprintln!("  Window func extra:{rc_window_extra:6}");
    eprintln!("  Col alias mismatch:{rc_col_alias:6}");
    eprintln!("  Missing sources:  {rc_missing_src:6}");
    eprintln!("  Extra sources:    {rc_extra_src:6}");
    eprintln!("  Other:            {rc_other:6}");

    let fail_total = rc_empty
        + rc_unnamed_col
        + rc_star_expand
        + rc_union_branch
        + rc_window_extra
        + rc_col_alias
        + rc_missing_src
        + rc_extra_src
        + rc_other;
    eprintln!("  ---");
    eprintln!("  Total failures:   {fail_total:6}");
}
