//! Detailed failure categorization for Snowflake ground truth lineage.
//!
//! Classifies every non-exact-match query into one of four buckets:
//!   - only_extra_output_cols  — we produce extra output columns, but every
//!     expected column matches exactly (harmless surplus)
//!   - only_missing_output_cols — expected has columns we don't emit at all
//!   - wrong_sources — same output columns, but source lists disagree
//!   - mixed — combination of the above
//!
//! Within `wrong_sources`, further breaks down into superset / subset / disjoint.
//! Within `only_extra_output_cols`, tallies the extra column names by frequency.
//!
//! Run:  cargo run --example gt_failure_analysis --release

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

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

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

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

/// Sort + dedup a column-lineage map for stable comparison.
fn normalize(map: &HashMap<String, Vec<String>>) -> HashMap<String, Vec<String>> {
    map.iter()
        .map(|(k, v)| {
            let mut sorted = v.clone();
            sorted.sort();
            sorted.dedup();
            (k.clone(), sorted)
        })
        .collect()
}

// ---------------------------------------------------------------------------
// Source-list comparison outcome for a single output column
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SourceCmp {
    Exact,
    Superset, // actual is a strict superset of expected
    Subset,   // actual is a strict subset of expected
    Disjoint, // no overlap at all
    Overlap,  // partial overlap, neither pure superset nor subset
}

fn compare_sources(actual: &[String], expected: &[String]) -> SourceCmp {
    if actual == expected {
        return SourceCmp::Exact;
    }
    let act_present_in_exp = expected.iter().all(|e| actual.contains(e));
    let exp_present_in_act = actual.iter().all(|a| expected.contains(a));

    match (act_present_in_exp, exp_present_in_act) {
        (true, false) => SourceCmp::Superset,
        (false, true) => SourceCmp::Subset,
        _ => {
            // Check for zero overlap
            let any_overlap = actual.iter().any(|a| expected.contains(a));
            if any_overlap {
                SourceCmp::Overlap
            } else {
                SourceCmp::Disjoint
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Per-query failure category
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum FailureKind {
    ExactMatch,
    ParseFail,
    EmptyActual,
    OnlyExtraOutputCols,
    OnlyMissingOutputCols,
    WrongSources,
    Mixed,
}

impl std::fmt::Display for FailureKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ExactMatch => write!(f, "exact_match"),
            Self::ParseFail => write!(f, "parse_fail"),
            Self::EmptyActual => write!(f, "empty_actual"),
            Self::OnlyExtraOutputCols => write!(f, "only_extra_output_cols"),
            Self::OnlyMissingOutputCols => write!(f, "only_missing_output_cols"),
            Self::WrongSources => write!(f, "wrong_sources"),
            Self::Mixed => write!(f, "mixed"),
        }
    }
}

// ---------------------------------------------------------------------------
// Counters
// ---------------------------------------------------------------------------

#[derive(Default)]
struct Stats {
    /// Per failure-kind query counts.
    kind_counts: HashMap<FailureKind, u64>,

    // -- only_extra_output_cols detail --
    extra_col_freq: HashMap<String, u64>,

    // -- wrong_sources detail --
    ws_superset: u64,
    ws_subset: u64,
    ws_disjoint: u64,
    ws_overlap: u64,

    // -- per-column totals for wrong_sources --
    ws_col_superset: u64,
    ws_col_subset: u64,
    ws_col_disjoint: u64,
    ws_col_overlap: u64,
    ws_col_exact: u64,
}

impl Stats {
    fn inc(&mut self, kind: FailureKind) {
        *self.kind_counts.entry(kind).or_default() += 1;
    }

    fn total(&self) -> u64 {
        self.kind_counts.values().sum()
    }

    fn count(&self, kind: FailureKind) -> u64 {
        *self.kind_counts.get(&kind).unwrap_or(&0)
    }

    fn pct(&self, n: u64) -> f64 {
        let t = self.total();
        if t == 0 {
            0.0
        } else {
            n as f64 / t as f64 * 100.0
        }
    }
}

// ---------------------------------------------------------------------------
// 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 stats = Stats::default();

    // Collect a few sample cases per category for the report
    let sample_limit = 5usize;
    let mut samples_wrong_sources: Vec<(
        String,
        HashMap<String, Vec<String>>,
        HashMap<String, Vec<String>>,
    )> = Vec::new();
    let mut samples_missing_cols: Vec<(String, Vec<String>)> = Vec::new();

    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,
        };

        let schema = schema_from_json(&case.schema);
        let expected = normalize(&case.expected);

        // Run lineage engine (projection mode — no options)
        let result = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => r,
            Err(_) => {
                stats.inc(FailureKind::ParseFail);
                continue;
            }
        };

        let actual = normalize(&result.column_dict);

        // -- Fast path: exact match --
        if actual == expected {
            stats.inc(FailureKind::ExactMatch);
            continue;
        }

        // -- Empty actual --
        if actual.is_empty() && !expected.is_empty() {
            stats.inc(FailureKind::EmptyActual);
            continue;
        }

        // -- Classify the mismatch --
        let extra_cols: Vec<String> = actual
            .keys()
            .filter(|k| !expected.contains_key(*k))
            .cloned()
            .collect();
        let missing_cols: Vec<String> = expected
            .keys()
            .filter(|k| !actual.contains_key(*k))
            .cloned()
            .collect();
        let common_cols: Vec<&String> = expected
            .keys()
            .filter(|k| actual.contains_key(*k))
            .collect();

        // Do all common columns match exactly on sources?
        let all_common_exact = common_cols
            .iter()
            .all(|col| actual.get(*col) == expected.get(*col));

        let has_extra = !extra_cols.is_empty();
        let has_missing = !missing_cols.is_empty();
        let has_wrong = !all_common_exact;

        let kind = match (has_extra, has_missing, has_wrong) {
            (true, false, false) => FailureKind::OnlyExtraOutputCols,
            (false, true, false) => FailureKind::OnlyMissingOutputCols,
            (false, false, true) => FailureKind::WrongSources,
            // extra cols + all common exact + no missing = still OnlyExtraOutputCols
            // (already covered above)
            // missing cols + all common exact + no extra = still OnlyMissingOutputCols
            // (already covered above)
            // wrong sources with NO extra/missing columns = WrongSources
            // (already covered above)
            _ => {
                // Any combination with wrong sources + extra/missing = Mixed
                // Unless it's just extra+missing but all common are exact
                if !has_wrong && has_extra && has_missing {
                    FailureKind::Mixed // extra + missing but no source mismatch
                } else {
                    FailureKind::Mixed
                }
            }
        };

        stats.inc(kind);

        // -- Drill into category-specific details --
        match kind {
            FailureKind::OnlyExtraOutputCols => {
                for col in &extra_cols {
                    *stats.extra_col_freq.entry(col.clone()).or_default() += 1;
                }
            }
            FailureKind::OnlyMissingOutputCols => {
                if samples_missing_cols.len() < sample_limit {
                    samples_missing_cols.push((case.name.clone(), missing_cols.clone()));
                }
            }
            FailureKind::WrongSources => {
                // Per-column breakdown within wrong_sources queries
                let mut q_has_superset = false;
                let mut q_has_subset = false;
                let mut q_has_disjoint = false;
                let mut q_has_overlap = false;

                for col in &common_cols {
                    let act = actual.get(*col).unwrap();
                    let exp = expected.get(*col).unwrap();
                    match compare_sources(act, exp) {
                        SourceCmp::Exact => {
                            stats.ws_col_exact += 1;
                        }
                        SourceCmp::Superset => {
                            stats.ws_col_superset += 1;
                            q_has_superset = true;
                        }
                        SourceCmp::Subset => {
                            stats.ws_col_subset += 1;
                            q_has_subset = true;
                        }
                        SourceCmp::Disjoint => {
                            stats.ws_col_disjoint += 1;
                            q_has_disjoint = true;
                        }
                        SourceCmp::Overlap => {
                            stats.ws_col_overlap += 1;
                            q_has_overlap = true;
                        }
                    }
                }

                // Classify the query-level wrong_sources flavour by the
                // dominant column-level pattern.
                if q_has_superset && !q_has_subset && !q_has_disjoint && !q_has_overlap {
                    stats.ws_superset += 1;
                } else if q_has_subset && !q_has_superset && !q_has_disjoint && !q_has_overlap {
                    stats.ws_subset += 1;
                } else if q_has_disjoint && !q_has_superset && !q_has_subset && !q_has_overlap {
                    stats.ws_disjoint += 1;
                } else {
                    stats.ws_overlap += 1; // mixed within wrong_sources
                }

                if samples_wrong_sources.len() < sample_limit {
                    samples_wrong_sources.push((
                        case.name.clone(),
                        actual.clone(),
                        expected.clone(),
                    ));
                }
            }
            FailureKind::Mixed => {
                // Count extra col frequencies even for mixed
                for col in &extra_cols {
                    *stats.extra_col_freq.entry(col.clone()).or_default() += 1;
                }
            }
            _ => {}
        }

        // Progress
        if (line_idx + 1) % 2000 == 0 {
            eprint!("\r  processed {} lines ...", line_idx + 1);
        }
    }
    eprintln!("\r  done.                        ");

    // -----------------------------------------------------------------------
    // Report
    // -----------------------------------------------------------------------

    let total = stats.total();
    println!();
    println!("========================================================");
    println!("  Snowflake Ground Truth — Failure Categorization");
    println!("========================================================");
    println!();
    println!("  Total queries:               {:>7}", total);
    println!();
    println!(
        "  exact_match:                 {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::ExactMatch),
        stats.pct(stats.count(FailureKind::ExactMatch))
    );
    println!(
        "  parse_fail:                  {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::ParseFail),
        stats.pct(stats.count(FailureKind::ParseFail))
    );
    println!(
        "  empty_actual:                {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::EmptyActual),
        stats.pct(stats.count(FailureKind::EmptyActual))
    );
    println!("  -----------------------------------------------");
    println!(
        "  only_extra_output_cols:       {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::OnlyExtraOutputCols),
        stats.pct(stats.count(FailureKind::OnlyExtraOutputCols))
    );
    println!(
        "  only_missing_output_cols:     {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::OnlyMissingOutputCols),
        stats.pct(stats.count(FailureKind::OnlyMissingOutputCols))
    );
    println!(
        "  wrong_sources:               {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::WrongSources),
        stats.pct(stats.count(FailureKind::WrongSources))
    );
    println!(
        "  mixed:                       {:>7}  ({:>5.1}%)",
        stats.count(FailureKind::Mixed),
        stats.pct(stats.count(FailureKind::Mixed))
    );

    // -- wrong_sources detail --
    let ws = stats.count(FailureKind::WrongSources);
    if ws > 0 {
        println!();
        println!("  --- wrong_sources breakdown (query-level) ---");
        println!(
            "    pure superset (extra sources only):  {:>6}",
            stats.ws_superset
        );
        println!(
            "    pure subset (missing sources only):  {:>6}",
            stats.ws_subset
        );
        println!(
            "    pure disjoint (zero overlap):        {:>6}",
            stats.ws_disjoint
        );
        println!(
            "    mixed overlap:                       {:>6}",
            stats.ws_overlap
        );

        let ws_cols = stats.ws_col_exact
            + stats.ws_col_superset
            + stats.ws_col_subset
            + stats.ws_col_disjoint
            + stats.ws_col_overlap;
        if ws_cols > 0 {
            println!();
            println!(
                "  --- wrong_sources breakdown (per-column, {} cols) ---",
                ws_cols
            );
            println!(
                "    exact:    {:>6}  ({:>5.1}%)",
                stats.ws_col_exact,
                stats.ws_col_exact as f64 / ws_cols as f64 * 100.0
            );
            println!(
                "    superset: {:>6}  ({:>5.1}%)",
                stats.ws_col_superset,
                stats.ws_col_superset as f64 / ws_cols as f64 * 100.0
            );
            println!(
                "    subset:   {:>6}  ({:>5.1}%)",
                stats.ws_col_subset,
                stats.ws_col_subset as f64 / ws_cols as f64 * 100.0
            );
            println!(
                "    disjoint: {:>6}  ({:>5.1}%)",
                stats.ws_col_disjoint,
                stats.ws_col_disjoint as f64 / ws_cols as f64 * 100.0
            );
            println!(
                "    overlap:  {:>6}  ({:>5.1}%)",
                stats.ws_col_overlap,
                stats.ws_col_overlap as f64 / ws_cols as f64 * 100.0
            );
        }
    }

    // -- extra output column frequency --
    if !stats.extra_col_freq.is_empty() {
        let mut freq: Vec<_> = stats.extra_col_freq.iter().collect();
        freq.sort_by(|a, b| b.1.cmp(a.1));
        let total_extra: u64 = freq.iter().map(|(_, c)| **c).sum();
        let unique = freq.len();
        println!();
        println!("  --- extra output column frequencies ---");
        println!("    total extra column occurrences: {total_extra}");
        println!("    unique extra column names:      {unique}");
        println!();
        println!("    top 30:");
        for (col, count) in freq.iter().take(30) {
            println!("      {:<50} {:>6}", col, count);
        }
        if unique > 30 {
            let tail: u64 = freq.iter().skip(30).map(|(_, c)| **c).sum();
            println!(
                "      {:<50} {:>6}",
                format!("... ({} more)", unique - 30),
                tail
            );
        }
    }

    // -- sample wrong_sources cases --
    if !samples_wrong_sources.is_empty() {
        println!();
        println!("  --- sample wrong_sources cases ---");
        for (name, actual, expected) in &samples_wrong_sources {
            println!();
            println!("    case: {name}");
            let all_cols: Vec<&String> = {
                let mut s: Vec<&String> = actual.keys().chain(expected.keys()).collect();
                s.sort();
                s.dedup();
                s
            };
            for col in all_cols {
                let act = actual.get(col);
                let exp = expected.get(col);
                match (act, exp) {
                    (Some(a), Some(e)) if a != e => {
                        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: {missing:?}");
                        }
                        if !extra.is_empty() {
                            println!("        extra:   {extra:?}");
                        }
                    }
                    (None, Some(_)) => {
                        println!("      col {col}: MISSING in actual");
                    }
                    (Some(_), None) => {
                        println!("      col {col}: EXTRA in actual");
                    }
                    _ => {}
                }
            }
        }
    }

    // -- sample missing output col cases --
    if !samples_missing_cols.is_empty() {
        println!();
        println!("  --- sample only_missing_output_cols cases ---");
        for (name, cols) in &samples_missing_cols {
            println!("    {name}: missing cols = {cols:?}");
        }
    }

    println!();
    println!("========================================================");
}
