//! Categorize all failing ground truth cases and print a summary table.
//!
//! Usage:
//!   cargo run --example categorize --release

use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use serde::Deserialize;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::io::BufRead;
use std::path::Path;

#[derive(Debug, Deserialize)]
struct Case {
    name: String,
    sql: String,
    schema: serde_json::Value,
    expected: HashMap<String, Vec<String>>,
}

fn schema_from_json(value: &serde_json::Value) -> HashMap<String, serde_json::Value> {
    match value {
        serde_json::Value::Object(map) => map.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
        _ => HashMap::new(),
    }
}

fn normalize(map: &HashMap<String, Vec<String>>) -> HashMap<String, BTreeSet<String>> {
    map.iter()
        .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
        .collect()
}

fn classify(
    actual: &HashMap<String, BTreeSet<String>>,
    expected: &HashMap<String, BTreeSet<String>>,
) -> &'static str {
    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, |a| {
            let e = expected.get(k).unwrap();
            a != e
        })
    });

    match (has_extra_cols, has_missing_cols, has_wrong_sources) {
        (true, false, false) => "only_extra_output_cols",
        (false, true, false) => "only_missing_output_cols",
        (false, false, true) => "wrong_sources",
        (true, false, true) => "extra_cols+wrong_sources",
        (false, true, true) => "missing_cols+wrong_sources",
        (true, true, _) => "mixed",
        _ => "unknown",
    }
}

fn main() {
    let fixture_path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../../tests/fixtures/lineage/snowflake_ground_truth.jsonl");
    let file = std::fs::File::open(&fixture_path).expect("open fixture");
    let reader = std::io::BufReader::with_capacity(1024 * 1024, file);
    let dialect = SqlDialect::Snowflake;

    let mut total = 0usize;
    let mut exact_match = 0usize;
    let mut parse_fail = 0usize;
    let mut empty_actual = 0usize;

    // Category -> (count, first 3 example names)
    let mut categories: BTreeMap<&'static str, (usize, Vec<String>)> = BTreeMap::new();
    // Sub-classification for wrong_sources: superset vs subset vs disjoint
    let mut wrong_src_sub: BTreeMap<&'static str, usize> = BTreeMap::new();

    let start = std::time::Instant::now();

    for line in reader.lines() {
        let line = line.unwrap();
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let case: Case = match serde_json::from_str(line) {
            Ok(c) => c,
            Err(_) => continue,
        };

        total += 1;

        let schema = schema_from_json(&case.schema);
        let result = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => r,
            Err(_e) => {
                parse_fail += 1;
                continue;
            }
        };

        let actual = normalize(&result.column_dict);
        let expected = normalize(&case.expected);

        if actual == expected {
            exact_match += 1;
            continue;
        }

        // Check empty actual
        if actual.is_empty() || actual.values().all(|s| s.is_empty()) {
            empty_actual += 1;
        }

        let category = classify(&actual, &expected);
        let entry = categories
            .entry(category)
            .or_insert_with(|| (0, Vec::new()));
        entry.0 += 1;
        if entry.1.len() < 3 {
            entry.1.push(case.name.clone());
        }

        // Sub-classify wrong_sources cases: for each shared column, check
        // if actual is superset, subset, or disjoint relative to expected
        if category == "wrong_sources" || category.contains("wrong_sources") {
            let mut any_superset = false;
            let mut any_subset = false;
            let mut any_disjoint = false;
            for key in expected.keys() {
                if let Some(act) = actual.get(key) {
                    let exp = expected.get(key).unwrap();
                    if act == exp {
                        continue;
                    }
                    let act_superset = exp.is_subset(act);
                    let act_subset = act.is_subset(exp);
                    if act_superset && !act_subset {
                        any_superset = true;
                    } else if act_subset && !act_superset {
                        any_subset = true;
                    } else if !act_superset && !act_subset {
                        any_disjoint = true;
                    }
                }
            }
            let sub = if any_disjoint {
                "disjoint_sources"
            } else if any_superset && any_subset {
                "mixed_super_sub"
            } else if any_superset {
                "superset_only"
            } else if any_subset {
                "subset_only"
            } else {
                "other"
            };
            *wrong_src_sub.entry(sub).or_insert(0) += 1;
        }
    }

    let elapsed = start.elapsed();
    let fail_count: usize = categories.values().map(|(c, _)| c).sum();

    // Print summary
    println!("╔══════════════════════════════════════════════════════════════════╗");
    println!("║            SNOWFLAKE GROUND TRUTH CATEGORIZATION                ║");
    println!("╠══════════════════════════════════════════════════════════════════╣");
    println!(
        "║  Total cases:       {:>6}                                      ║",
        total
    );
    println!(
        "║  Exact match:       {:>6}  ({:>5.1}%)                            ║",
        exact_match,
        exact_match as f64 / total as f64 * 100.0
    );
    println!(
        "║  Parse fail:        {:>6}  ({:>5.1}%)                            ║",
        parse_fail,
        parse_fail as f64 / total as f64 * 100.0
    );
    println!(
        "║  Empty actual:      {:>6}  ({:>5.1}%)                            ║",
        empty_actual,
        empty_actual as f64 / total as f64 * 100.0
    );
    println!(
        "║  Total failing:     {:>6}  ({:>5.1}%)                            ║",
        fail_count,
        fail_count as f64 / total as f64 * 100.0
    );
    println!(
        "║  Elapsed:           {:>6.1}s                                     ║",
        elapsed.as_secs_f64()
    );
    println!("╚══════════════════════════════════════════════════════════════════╝");

    println!();
    println!("┌──────────────────────────────┬────────┬────────┐");
    println!("│ Category                     │  Count │      % │");
    println!("├──────────────────────────────┼────────┼────────┤");
    // Sort by count descending
    let mut sorted: Vec<_> = categories.iter().collect();
    sorted.sort_by(|a, b| b.1.0.cmp(&a.1.0));
    for (cat, (count, _examples)) in &sorted {
        println!(
            "│ {:<28} │ {:>6} │ {:>5.1}% │",
            cat,
            count,
            *count as f64 / total as f64 * 100.0
        );
    }
    println!("└──────────────────────────────┴────────┴────────┘");

    // Wrong-sources sub-classification
    if !wrong_src_sub.is_empty() {
        println!();
        println!("Wrong-source sub-classification (across all categories with wrong sources):");
        println!("┌──────────────────────────────┬────────┐");
        println!("│ Sub-category                 │  Count │");
        println!("├──────────────────────────────┼────────┤");
        let mut sorted_sub: Vec<_> = wrong_src_sub.iter().collect();
        sorted_sub.sort_by(|a, b| b.1.cmp(a.1));
        for (sub, count) in &sorted_sub {
            println!("│ {:<28} │ {:>6} │", sub, count);
        }
        println!("└──────────────────────────────┴────────┘");
    }

    // Print examples for each category
    println!();
    println!("═══ EXAMPLES PER CATEGORY (first 3) ═══════════════════════════════");
    for (cat, (_count, examples)) in &sorted {
        println!("\n  {} ({} cases):", cat, _count);
        for (i, name) in examples.iter().enumerate() {
            println!("    {}. {}", i + 1, name);
        }
    }

    println!();
    println!("Done. Run eyeball on specific cases:");
    println!("  cargo run --example eyeball --release -- <case_name>");
    println!("  cargo run --example eyeball --release -- --category <category_name>");
}
