//! Output column analysis for extra and missing columns in lineage results.
//!
//! Two-part analysis:
//!
//! **Part 1 — Extra output columns (`only_extra_output_cols`):**
//!   Tallies which extra column names appear most frequently, then simulates
//!   filtering those columns from our output to see how many queries would
//!   become exact matches.
//!
//! **Part 2 — Missing output columns (`only_missing_output_cols`):**
//!   Tallies the most common missing column names, categorises root causes,
//!   and prints 10 sample cases with full SQL / expected / actual detail.
//!
//! Run:  cargo run --release --example output_col_analysis

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

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

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

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

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

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

/// All column names defined in the schema (upper-cased for comparison).
fn schema_columns(schema: &serde_json::Value) -> HashSet<String> {
    let mut cols = HashSet::new();
    if let serde_json::Value::Object(tables) = schema {
        for (_table, columns) in tables {
            if let serde_json::Value::Object(col_map) = columns {
                for col_name in col_map.keys() {
                    cols.insert(col_name.to_uppercase());
                }
            }
        }
    }
    cols
}

/// Detect SQL features via simple substring matching.
fn has_select_star(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    upper.contains("SELECT *") || upper.contains("SELECT  *")
}

fn has_alias_dot_star(sql: &str) -> bool {
    // e.g. `a.*`, `src.*`
    let upper = sql.to_uppercase();
    // Look for pattern: word.* not preceded by /*
    for (i, _) in upper.match_indices(".*") {
        if i > 0 {
            let prev_char = upper.as_bytes()[i - 1];
            if prev_char.is_ascii_alphanumeric() || prev_char == b'_' {
                return true;
            }
        }
    }
    false
}

fn has_cte(sql: &str) -> bool {
    sql.to_uppercase().trim_start().starts_with("WITH ")
}

fn has_union(sql: &str) -> bool {
    sql.to_uppercase().contains("UNION")
}

fn has_lateral_flatten(sql: &str) -> bool {
    sql.to_uppercase().contains("LATERAL FLATTEN") || sql.to_uppercase().contains("TABLE(FLATTEN")
}

// ---------------------------------------------------------------------------
// Category classifiers
// ---------------------------------------------------------------------------

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

fn classify(
    actual: &HashMap<String, Vec<String>>,
    expected: &HashMap<String, Vec<String>>,
) -> QueryCategory {
    if actual == expected {
        return QueryCategory::ExactMatch;
    }
    if actual.is_empty() && !expected.is_empty() {
        return QueryCategory::EmptyActual;
    }

    let has_extra = actual.keys().any(|k| !expected.contains_key(k));
    let has_missing = expected.keys().any(|k| !actual.contains_key(k));
    let has_wrong = expected.keys().any(|k| {
        actual
            .get(k)
            .map_or(false, |act| act != expected.get(k).unwrap())
    });

    match (has_extra, has_missing, has_wrong) {
        (true, false, false) => QueryCategory::OnlyExtraOutputCols,
        (false, true, false) => QueryCategory::OnlyMissingOutputCols,
        (false, false, true) => QueryCategory::WrongSources,
        _ => QueryCategory::Mixed,
    }
}

// ---------------------------------------------------------------------------
// Sample collector for missing output cols
// ---------------------------------------------------------------------------

struct MissingSample {
    name: String,
    sql: String,
    expected_cols: Vec<String>,
    actual_cols: Vec<String>,
    missing_cols: Vec<String>,
    has_star: bool,
    has_alias_star: bool,
    has_cte: bool,
    has_union: bool,
    has_flatten: bool,
}

// ---------------------------------------------------------------------------
// 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());
        std::process::exit(1);
    }

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

    // --- Counters ---
    let mut total = 0u64;
    let mut exact = 0u64;

    // --- Extra cols ---
    let mut extra_cases = 0u64;
    let mut extra_col_freq: HashMap<String, u64> = HashMap::new();
    // For each extra case, record which extra col names it has
    let mut extra_case_sets: Vec<HashSet<String>> = Vec::new();

    // --- Missing cols ---
    let mut missing_cases = 0u64;
    let mut missing_col_freq: HashMap<String, u64> = HashMap::new();
    let mut missing_col_cause: HashMap<String, HashMap<String, u64>> = HashMap::new(); // col -> cause -> count
    let mut missing_samples: Vec<MissingSample> = Vec::new();

    // --- Missing cols SQL feature counts ---
    let mut miss_has_star = 0u64;
    let mut miss_has_alias_star = 0u64;
    let mut miss_has_cte_count = 0u64;
    let mut miss_has_union_count = 0u64;
    let mut miss_has_flatten_count = 0u64;

    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_map(&case.schema);
        let expected = normalize(&case.expected);

        let result = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => r,
            Err(_) => continue,
        };

        let actual = normalize(&result.column_dict);
        let category = classify(&actual, &expected);

        if category == QueryCategory::ExactMatch {
            exact += 1;
            continue;
        }

        // ---- EXTRA OUTPUT COLS ----
        if category == QueryCategory::OnlyExtraOutputCols {
            extra_cases += 1;
            let extras: HashSet<String> = actual
                .keys()
                .filter(|k| !expected.contains_key(*k))
                .cloned()
                .collect();
            for col in &extras {
                *extra_col_freq.entry(col.clone()).or_default() += 1;
            }
            extra_case_sets.push(extras);
        }

        // ---- MISSING OUTPUT COLS ----
        if category == QueryCategory::OnlyMissingOutputCols {
            missing_cases += 1;
            let schema_cols = schema_columns(&case.schema);
            let missing: Vec<String> = expected
                .keys()
                .filter(|k| !actual.contains_key(*k))
                .cloned()
                .collect();

            let star = has_select_star(&case.sql);
            let alias_star = has_alias_dot_star(&case.sql);
            let cte = has_cte(&case.sql);
            let union_q = has_union(&case.sql);
            let flatten = has_lateral_flatten(&case.sql);

            if star {
                miss_has_star += 1;
            }
            if alias_star {
                miss_has_alias_star += 1;
            }
            if cte {
                miss_has_cte_count += 1;
            }
            if union_q {
                miss_has_union_count += 1;
            }
            if flatten {
                miss_has_flatten_count += 1;
            }

            for col in &missing {
                *missing_col_freq.entry(col.clone()).or_default() += 1;

                // Classify cause
                let cause = if schema_cols.contains(&col.to_uppercase()) && (star || alias_star) {
                    "star_expansion"
                } else if schema_cols.contains(&col.to_uppercase()) {
                    "schema_col_missed"
                } else {
                    "expression_alias"
                };
                *missing_col_cause
                    .entry(col.clone())
                    .or_default()
                    .entry(cause.to_string())
                    .or_default() += 1;
            }

            if missing_samples.len() < 10 {
                let mut exp_cols: Vec<String> = expected.keys().cloned().collect();
                exp_cols.sort();
                let mut act_cols: Vec<String> = actual.keys().cloned().collect();
                act_cols.sort();
                missing_samples.push(MissingSample {
                    name: case.name.clone(),
                    sql: case.sql.clone(),
                    expected_cols: exp_cols,
                    actual_cols: act_cols,
                    missing_cols: missing.clone(),
                    has_star: star,
                    has_alias_star: alias_star,
                    has_cte: cte,
                    has_union: union_q,
                    has_flatten: flatten,
                });
            }
        }

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

    // =======================================================================
    // REPORT — PART 1: Extra Output Columns
    // =======================================================================

    println!();
    println!("================================================================");
    println!("  PART 1: Extra Output Column Analysis");
    println!("================================================================");
    println!();
    println!("  Total queries:     {:>7}", total);
    println!(
        "  Exact matches:     {:>7}  ({:.1}%)",
        exact,
        exact as f64 / total as f64 * 100.0
    );
    println!(
        "  only_extra_output: {:>7}  ({:.1}%)",
        extra_cases,
        extra_cases as f64 / total as f64 * 100.0
    );
    println!();

    // Top extra column names by frequency
    let mut freq_vec: Vec<_> = extra_col_freq.iter().collect();
    freq_vec.sort_by(|a, b| b.1.cmp(a.1));

    println!("  Top 30 extra column names:");
    println!("  {:<55} {:>6} {:>7}", "Column", "Count", "% of cases");
    println!("  {}", "-".repeat(70));
    for (col, count) in freq_vec.iter().take(30) {
        println!(
            "  {:<55} {:>6} {:>6.1}%",
            col,
            count,
            **count as f64 / extra_cases as f64 * 100.0
        );
    }

    // --- Filterable column simulation ---
    // Build a set of candidate filter columns (those appearing in >=1% of extra cases)
    let min_threshold = (extra_cases as f64 * 0.01).ceil() as u64;
    let candidate_filter: HashSet<String> = extra_col_freq
        .iter()
        .filter(|(_, count)| **count >= min_threshold)
        .map(|(col, _)| col.clone())
        .collect();

    println!();
    println!("  --- Filter simulation ---");
    println!();
    println!(
        "  Candidate filter columns (appearing in >= {:.0}% = {} cases):",
        1.0, min_threshold
    );
    let mut candidates_sorted: Vec<_> = candidate_filter.iter().cloned().collect();
    candidates_sorted.sort();
    for c in &candidates_sorted {
        println!("    {}", c);
    }

    // For each extra case, would removing candidate columns make it exact?
    let mut would_become_exact = 0u64;
    let mut would_become_exact_with_specific: HashMap<String, u64> = HashMap::new();

    for extras in &extra_case_sets {
        let remaining: HashSet<String> = extras
            .iter()
            .filter(|c| !candidate_filter.contains(*c))
            .cloned()
            .collect();
        if remaining.is_empty() {
            would_become_exact += 1;
            // Credit each candidate that was present
            for c in extras {
                if candidate_filter.contains(c) {
                    *would_become_exact_with_specific
                        .entry(c.clone())
                        .or_default() += 1;
                }
            }
        }
    }

    println!();
    println!(
        "  If we filter out ALL {} candidate columns from our output:",
        candidate_filter.len()
    );
    println!(
        "    {} / {} extra-only cases would become exact matches ({:.1}%)",
        would_become_exact,
        extra_cases,
        would_become_exact as f64 / extra_cases as f64 * 100.0
    );
    println!(
        "    New exact match total: {} / {} ({:.1}%)",
        exact + would_become_exact,
        total,
        (exact + would_become_exact) as f64 / total as f64 * 100.0
    );

    // Now try the top-N most common columns specifically
    println!();
    println!("  Incremental filtering (cumulative, top extra columns by frequency):");
    println!(
        "  {:<55} {:>10} {:>10}",
        "Filter set", "Recovered", "New total %"
    );
    println!("  {}", "-".repeat(77));

    let top_n_cols: Vec<String> = freq_vec
        .iter()
        .take(20)
        .map(|(c, _)| c.to_string())
        .collect();
    let mut cumulative_filter: HashSet<String> = HashSet::new();

    for col in &top_n_cols {
        cumulative_filter.insert(col.clone());

        let mut recovered = 0u64;
        for extras in &extra_case_sets {
            let remaining: HashSet<String> = extras
                .iter()
                .filter(|c| !cumulative_filter.contains(*c))
                .cloned()
                .collect();
            if remaining.is_empty() {
                recovered += 1;
            }
        }

        let new_pct = (exact + recovered) as f64 / total as f64 * 100.0;
        println!("  + {:<52} {:>10} {:>9.1}%", col, recovered, new_pct);
    }

    // =======================================================================
    // REPORT — PART 2: Missing Output Columns
    // =======================================================================

    println!();
    println!("================================================================");
    println!("  PART 2: Missing Output Column Analysis");
    println!("================================================================");
    println!();
    println!(
        "  only_missing_output: {:>7}  ({:.1}%)",
        missing_cases,
        missing_cases as f64 / total as f64 * 100.0
    );
    println!();

    println!("  SQL feature prevalence in missing-col queries:");
    println!(
        "    SELECT *:         {:>5}  ({:.1}%)",
        miss_has_star,
        miss_has_star as f64 / missing_cases.max(1) as f64 * 100.0
    );
    println!(
        "    alias.*:          {:>5}  ({:.1}%)",
        miss_has_alias_star,
        miss_has_alias_star as f64 / missing_cases.max(1) as f64 * 100.0
    );
    println!(
        "    CTE (WITH):       {:>5}  ({:.1}%)",
        miss_has_cte_count,
        miss_has_cte_count as f64 / missing_cases.max(1) as f64 * 100.0
    );
    println!(
        "    UNION:            {:>5}  ({:.1}%)",
        miss_has_union_count,
        miss_has_union_count as f64 / missing_cases.max(1) as f64 * 100.0
    );
    println!(
        "    LATERAL FLATTEN:  {:>5}  ({:.1}%)",
        miss_has_flatten_count,
        miss_has_flatten_count as f64 / missing_cases.max(1) as f64 * 100.0
    );

    // Top missing column names
    let mut miss_freq: Vec<_> = missing_col_freq.iter().collect();
    miss_freq.sort_by(|a, b| b.1.cmp(a.1));

    println!();
    println!("  Top 30 missing column names:");
    println!("  {:<40} {:>6} {:>30}", "Column", "Count", "Likely cause");
    println!("  {}", "-".repeat(78));
    for (col, count) in miss_freq.iter().take(30) {
        // Determine dominant cause
        let cause_str = if let Some(causes) = missing_col_cause.get(*col) {
            let mut cv: Vec<_> = causes.iter().collect();
            cv.sort_by(|a, b| b.1.cmp(a.1));
            cv.iter()
                .map(|(c, n)| format!("{}({})", c, n))
                .collect::<Vec<_>>()
                .join(", ")
        } else {
            "?".to_string()
        };
        println!("  {:<40} {:>6}  {}", col, count, cause_str);
    }

    // --- Samples ---
    if !missing_samples.is_empty() {
        println!();
        println!("  --- Sample missing-output-col cases (up to 10) ---");
        for (i, s) in missing_samples.iter().enumerate() {
            println!();
            println!("  [{}/{}] name: {}", i + 1, missing_samples.len(), s.name);
            println!(
                "       features: star={} alias_star={} cte={} union={} flatten={}",
                s.has_star, s.has_alias_star, s.has_cte, s.has_union, s.has_flatten
            );
            println!("       missing cols: {:?}", s.missing_cols);
            println!("       expected cols: {:?}", s.expected_cols);
            println!("       actual cols:   {:?}", s.actual_cols);
            // Print SQL (truncated to 500 chars)
            let sql_display = if s.sql.len() > 500 {
                format!("{}...", &s.sql[..500])
            } else {
                s.sql.clone()
            };
            println!(
                "       SQL: {}",
                sql_display.replace('\n', "\n            ")
            );
        }
    }

    // =======================================================================
    // SUMMARY
    // =======================================================================

    println!();
    println!("================================================================");
    println!("  SUMMARY");
    println!("================================================================");
    println!();
    println!(
        "  Current exact matches:  {} / {} ({:.1}%)",
        exact,
        total,
        exact as f64 / total as f64 * 100.0
    );

    // Best filter set: top 5 columns
    let top5: HashSet<String> = freq_vec
        .iter()
        .take(5)
        .map(|(c, _)| c.to_string())
        .collect();
    let mut top5_recovered = 0u64;
    for extras in &extra_case_sets {
        let remaining: HashSet<String> = extras
            .iter()
            .filter(|c| !top5.contains(*c))
            .cloned()
            .collect();
        if remaining.is_empty() {
            top5_recovered += 1;
        }
    }

    let top5_list: Vec<String> = freq_vec
        .iter()
        .take(5)
        .map(|(c, _)| c.to_string())
        .collect();
    println!();
    println!("  If we filter out top-5 extra columns from our output:");
    println!("    Columns: {:?}", top5_list);
    println!("    {} extra-only cases recovered", top5_recovered);
    println!(
        "    New exact match total: {} / {} ({:.1}%)",
        exact + top5_recovered,
        total,
        (exact + top5_recovered) as f64 / total as f64 * 100.0
    );
    println!();
    println!(
        "  Missing output columns affect {} queries ({:.1}%).",
        missing_cases,
        missing_cases as f64 / total as f64 * 100.0
    );
    println!("  These require engine fixes (not simple filtering).");
    println!();
    println!("================================================================");
}
