//! Categorize column lineage benchmark failures into actionable buckets.
//!
//! Reads the Snowflake ground truth JSONL, runs the lineage engine on each
//! case, and classifies the outcome into one of seven categories:
//!
//! | Category          | Meaning                                            |
//! |-------------------|----------------------------------------------------|
//! | exact_match       | Perfect match                                      |
//! | extra_cols_only   | All expected cols correct, but we emit extra ones   |
//! | missing_cols_only | We're missing output cols that expected has          |
//! | wrong_sources     | Same output cols, but source references differ      |
//! | mixed             | Combination of missing/extra cols AND source diffs  |
//! | empty_actual      | Engine produces no output at all                    |
//! | parse_error       | Engine returned an error                            |
//!
//! `wrong_sources` is further sub-categorized:
//!   - superset_sources:  we have ALL expected sources + extras
//!   - subset_sources:    expected has ALL our sources + more
//!   - disjoint_sources:  neither is a subset of the other
//!
//! Run:  cargo run --example failure_categories --release

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

// ---------------------------------------------------------------------------
// Ground truth case
// ---------------------------------------------------------------------------

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

// ---------------------------------------------------------------------------
// Failure categories
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Category {
    ExactMatch,
    ExtraColsOnly,
    MissingColsOnly,
    WrongSources,
    Mixed,
    EmptyActual,
    ParseError,
}

impl Category {
    fn label(self) -> &'static str {
        match self {
            Self::ExactMatch => "exact_match",
            Self::ExtraColsOnly => "extra_cols_only",
            Self::MissingColsOnly => "missing_cols_only",
            Self::WrongSources => "wrong_sources",
            Self::Mixed => "mixed",
            Self::EmptyActual => "empty_actual",
            Self::ParseError => "parse_error",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum WrongSourcesSub {
    Superset,
    Subset,
    Disjoint,
}

impl WrongSourcesSub {
    fn label(self) -> &'static str {
        match self {
            Self::Superset => "superset_sources",
            Self::Subset => "subset_sources",
            Self::Disjoint => "disjoint_sources",
        }
    }
}

// ---------------------------------------------------------------------------
// Sample case for display
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
struct SampleCase {
    name: String,
    sql_preview: String,
    expected: HashMap<String, Vec<String>>,
    actual: HashMap<String, Vec<String>>,
}

impl SampleCase {
    fn new(
        name: &str,
        sql: &str,
        expected: &HashMap<String, BTreeSet<String>>,
        actual: &HashMap<String, BTreeSet<String>>,
    ) -> Self {
        let preview = if sql.len() > 200 {
            format!("{}...", &sql[..200])
        } else {
            sql.to_string()
        };
        Self {
            name: name.to_string(),
            sql_preview: preview.replace('\n', " "),
            expected: expected
                .iter()
                .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
                .collect(),
            actual: actual
                .iter()
                .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
                .collect(),
        }
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

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(),
    }
}

/// Normalize a column_dict into sorted sets for stable comparison.
fn normalize(map: &HashMap<String, Vec<String>>) -> HashMap<String, BTreeSet<String>> {
    map.iter()
        .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
        .collect()
}

/// Classify source-list relationship for a single output column.
fn classify_sources(actual: &BTreeSet<String>, expected: &BTreeSet<String>) -> WrongSourcesSub {
    let exp_subset_of_act = expected.is_subset(actual);
    let act_subset_of_exp = actual.is_subset(expected);

    match (exp_subset_of_act, act_subset_of_exp) {
        (true, false) => WrongSourcesSub::Superset,
        (false, true) => WrongSourcesSub::Subset,
        _ => WrongSourcesSub::Disjoint,
    }
}

/// Classify a non-matching case into its primary category.
fn classify(
    actual: &HashMap<String, BTreeSet<String>>,
    expected: &HashMap<String, BTreeSet<String>>,
) -> Category {
    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| act != expected.get(k).unwrap())
    });

    match (has_extra_cols, has_missing_cols, has_wrong_sources) {
        (true, false, false) => Category::ExtraColsOnly,
        (false, true, false) => Category::MissingColsOnly,
        (false, false, true) => Category::WrongSources,
        _ => Category::Mixed,
    }
}

/// For a wrong_sources case, determine the dominant sub-category across
/// all shared columns with differing sources.
fn classify_wrong_sources_sub(
    actual: &HashMap<String, BTreeSet<String>>,
    expected: &HashMap<String, BTreeSet<String>>,
) -> WrongSourcesSub {
    let mut any_superset = false;
    let mut any_subset = false;
    let mut any_disjoint = false;

    for (col, exp_srcs) in expected {
        if let Some(act_srcs) = actual.get(col) {
            if act_srcs == exp_srcs {
                continue;
            }
            match classify_sources(act_srcs, exp_srcs) {
                WrongSourcesSub::Superset => any_superset = true,
                WrongSourcesSub::Subset => any_subset = true,
                WrongSourcesSub::Disjoint => any_disjoint = true,
            }
        }
    }

    // Dominant classification: disjoint wins over mixed super/sub
    if any_disjoint {
        WrongSourcesSub::Disjoint
    } else if any_superset && !any_subset {
        WrongSourcesSub::Superset
    } else if any_subset && !any_superset {
        WrongSourcesSub::Subset
    } else {
        // Both superset and subset columns present
        WrongSourcesSub::Disjoint
    }
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

fn main() {
    let fixture_path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../../tests/fixtures/lineage/snowflake_ground_truth.jsonl");

    if !fixture_path.exists() {
        eprintln!("Ground-truth fixture not found: {}", fixture_path.display());
        eprintln!("Expected at tests/fixtures/lineage/snowflake_ground_truth.jsonl");
        std::process::exit(1);
    }

    let file = std::fs::File::open(&fixture_path).expect("failed to open fixture");
    let reader = std::io::BufReader::with_capacity(1024 * 1024, file);
    let dialect = SqlDialect::Snowflake;

    let mut total = 0usize;
    let mut counts: HashMap<Category, usize> = HashMap::new();
    let mut ws_sub_counts: HashMap<WrongSourcesSub, usize> = HashMap::new();
    let mut samples: HashMap<Category, Vec<SampleCase>> = HashMap::new();
    let mut ws_sub_samples: HashMap<WrongSourcesSub, Vec<SampleCase>> = HashMap::new();

    let sample_limit = 3usize;
    let start = std::time::Instant::now();

    for (line_idx, line) in reader.lines().enumerate() {
        let line = line.expect("failed to read line");
        let line = line.trim().to_string();
        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 expected = normalize(&case.expected);

        let result = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => r,
            Err(_) => {
                *counts.entry(Category::ParseError).or_default() += 1;
                let entry = samples.entry(Category::ParseError).or_default();
                if entry.len() < sample_limit {
                    entry.push(SampleCase::new(
                        &case.name,
                        &case.sql,
                        &expected,
                        &HashMap::new(),
                    ));
                }
                continue;
            }
        };

        let actual = normalize(&result.column_dict);

        // Exact match
        if actual == expected {
            *counts.entry(Category::ExactMatch).or_default() += 1;
            continue;
        }

        // Empty actual
        if actual.is_empty() && !expected.is_empty() {
            *counts.entry(Category::EmptyActual).or_default() += 1;
            let entry = samples.entry(Category::EmptyActual).or_default();
            if entry.len() < sample_limit {
                entry.push(SampleCase::new(&case.name, &case.sql, &expected, &actual));
            }
            continue;
        }

        let category = classify(&actual, &expected);
        *counts.entry(category).or_default() += 1;

        // Collect samples for this category
        let entry = samples.entry(category).or_default();
        if entry.len() < sample_limit {
            entry.push(SampleCase::new(&case.name, &case.sql, &expected, &actual));
        }

        // Sub-classify wrong_sources
        if category == Category::WrongSources {
            let sub = classify_wrong_sources_sub(&actual, &expected);
            *ws_sub_counts.entry(sub).or_default() += 1;

            let sub_entry = ws_sub_samples.entry(sub).or_default();
            if sub_entry.len() < sample_limit {
                sub_entry.push(SampleCase::new(&case.name, &case.sql, &expected, &actual));
            }
        }

        if (line_idx + 1) % 2000 == 0 {
            eprint!("\r  processed {} / ~19875 ...", line_idx + 1);
        }
    }
    let elapsed = start.elapsed();
    eprintln!(
        "\r  done. ({:.1}s)                      ",
        elapsed.as_secs_f64()
    );

    // -----------------------------------------------------------------------
    // Summary table
    // -----------------------------------------------------------------------

    let pct = |n: usize| -> f64 {
        if total == 0 {
            0.0
        } else {
            n as f64 / total as f64 * 100.0
        }
    };

    let cat_count = |c: Category| -> usize { *counts.get(&c).unwrap_or(&0) };

    println!();
    println!("=================================================================");
    println!(
        "  FAILURE CATEGORIES — Snowflake Ground Truth ({} queries)",
        total
    );
    println!("=================================================================");
    println!();
    println!("  {:<22} {:>7} {:>8}", "Category", "Count", "%");
    println!("  {:-<22} {:-<7} {:-<8}", "", "", "");

    let display_order = [
        Category::ExactMatch,
        Category::ExtraColsOnly,
        Category::MissingColsOnly,
        Category::WrongSources,
        Category::Mixed,
        Category::EmptyActual,
        Category::ParseError,
    ];

    for cat in &display_order {
        let n = cat_count(*cat);
        println!("  {:<22} {:>7} {:>7.1}%", cat.label(), n, pct(n));
    }

    println!();
    println!(
        "  Failures total:        {:>7} ({:.1}%)",
        total - cat_count(Category::ExactMatch),
        pct(total - cat_count(Category::ExactMatch))
    );

    // -----------------------------------------------------------------------
    // Wrong sources sub-categories
    // -----------------------------------------------------------------------

    let ws_total = cat_count(Category::WrongSources);
    if ws_total > 0 {
        let ws_pct = |n: usize| -> f64 {
            if ws_total == 0 {
                0.0
            } else {
                n as f64 / ws_total as f64 * 100.0
            }
        };

        println!();
        println!(
            "  --- wrong_sources sub-categories ({} queries) ---",
            ws_total
        );
        println!();
        println!("  {:<22} {:>7} {:>8}", "Sub-category", "Count", "%");
        println!("  {:-<22} {:-<7} {:-<8}", "", "", "");

        let ws_subs = [
            WrongSourcesSub::Superset,
            WrongSourcesSub::Subset,
            WrongSourcesSub::Disjoint,
        ];
        for sub in &ws_subs {
            let n = *ws_sub_counts.get(sub).unwrap_or(&0);
            println!("  {:<22} {:>7} {:>7.1}%", sub.label(), n, ws_pct(n));
        }
    }

    // -----------------------------------------------------------------------
    // Sample cases per category
    // -----------------------------------------------------------------------

    println!();
    println!("=================================================================");
    println!("  SAMPLE CASES (up to 3 per category)");
    println!("=================================================================");

    for cat in &display_order {
        if *cat == Category::ExactMatch {
            continue;
        }
        let n = cat_count(*cat);
        if n == 0 {
            continue;
        }

        println!();
        println!("  ---- {} ({} cases) ----", cat.label(), n);

        if let Some(cases) = samples.get(cat) {
            for (i, sample) in cases.iter().enumerate() {
                println!();
                println!("    {}. name: {}", i + 1, sample.name);
                println!("       sql:  {}", sample.sql_preview);

                // Show expected vs actual column_dict
                let all_cols: BTreeSet<&String> =
                    sample.expected.keys().chain(sample.actual.keys()).collect();

                for col in &all_cols {
                    let exp = sample.expected.get(*col);
                    let act = sample.actual.get(*col);
                    match (exp, act) {
                        (Some(e), Some(a)) if e == a => {
                            // Match — skip for brevity
                        }
                        (Some(e), Some(a)) => {
                            println!("       col \"{col}\":");
                            println!("         expected: {e:?}");
                            println!("         actual:   {a:?}");
                        }
                        (Some(e), None) => {
                            println!("       col \"{col}\": MISSING in actual");
                            println!("         expected: {e:?}");
                        }
                        (None, Some(a)) => {
                            println!("       col \"{col}\": EXTRA in actual");
                            println!("         actual:   {a:?}");
                        }
                        (None, None) => {}
                    }
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Wrong sources sub-category samples
    // -----------------------------------------------------------------------

    if ws_total > 0 {
        println!();
        println!("=================================================================");
        println!("  WRONG_SOURCES SUB-CATEGORY SAMPLES (up to 3 each)");
        println!("=================================================================");

        let ws_subs = [
            WrongSourcesSub::Superset,
            WrongSourcesSub::Subset,
            WrongSourcesSub::Disjoint,
        ];
        for sub in &ws_subs {
            let n = *ws_sub_counts.get(sub).unwrap_or(&0);
            if n == 0 {
                continue;
            }

            println!();
            println!("  ---- {} ({} cases) ----", sub.label(), n);

            if let Some(cases) = ws_sub_samples.get(sub) {
                for (i, sample) in cases.iter().enumerate() {
                    println!();
                    println!("    {}. name: {}", i + 1, sample.name);
                    println!("       sql:  {}", sample.sql_preview);

                    let all_cols: BTreeSet<&String> =
                        sample.expected.keys().chain(sample.actual.keys()).collect();

                    for col in &all_cols {
                        let exp = sample.expected.get(*col);
                        let act = sample.actual.get(*col);
                        match (exp, act) {
                            (Some(e), Some(a)) if e != a => {
                                let missing: Vec<_> = e.iter().filter(|v| !a.contains(v)).collect();
                                let extra: Vec<_> = a.iter().filter(|v| !e.contains(v)).collect();
                                println!("       col \"{col}\":");
                                if !missing.is_empty() {
                                    println!("         missing sources: {missing:?}");
                                }
                                if !extra.is_empty() {
                                    println!("         extra sources:   {extra:?}");
                                }
                            }
                            _ => {}
                        }
                    }
                }
            }
        }
    }

    println!();
    println!("=================================================================");
    println!("  Done in {:.1}s", elapsed.as_secs_f64());
    println!("=================================================================");
}
