//! Diagnostic for "missing output column" — expected columns absent from actual
//! lineage when the engine *did* produce some output.
//!
//! Focuses on queries where `projection_column_lineage` returns a non-empty
//! `column_dict` but the expected ground truth contains columns not present in
//! the actual output.  Categorizes each missing column by root cause:
//!
//!   (a) Schema column we missed — column exists in at least one schema table
//!       but was not emitted as an output column.
//!   (b) Expression alias — column name does not appear in any schema table
//!       (e.g. computed aliases like `TOTAL_AMOUNT`, `RN`, `JSON_PAYLOAD`).
//!   (c) Star expansion failure — expected column matches a `*` pattern and the
//!       query contains SELECT * or alias.*.
//!
//! Additionally tracks per-query SQL features:
//!   - SELECT * / alias.* presence
//!   - CTE (WITH ... AS) presence
//!   - UNION / UNION ALL presence
//!
//! Run: cargo run --example missing_cols --release

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(Debug, Deserialize)]
struct Case {
    name: String,
    sql: String,
    schema: serde_json::Value,
    expected: HashMap<String, Vec<String>>,
}

/// Why an expected output column is missing from the actual lineage result.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum MissingCause {
    /// Column name exists in a schema table but was not emitted.
    SchemaColumnMissed,
    /// Column is a computed expression alias — not in any schema table.
    ExpressionAlias,
    /// Column exists in schema AND the query uses SELECT * / alias.* — likely
    /// a star expansion failure.
    StarExpansionFailure,
}

impl MissingCause {
    fn label(&self) -> &'static str {
        match self {
            MissingCause::SchemaColumnMissed => "(a) Schema column we missed",
            MissingCause::ExpressionAlias => "(b) Expression alias not in schema",
            MissingCause::StarExpansionFailure => "(c) Star expansion failure",
        }
    }
}

/// A single missing column with enough context for sample display.
#[derive(Debug, Clone)]
#[allow(dead_code)]
struct MissingEntry {
    query_name: String,
    sql_prefix: String,
    column: String,
    expected_sources: Vec<String>,
    actual_cols: Vec<String>,
    cause: MissingCause,
    has_star: bool,
    has_cte: bool,
    has_union: bool,
}

// ─────────────────────────────────────────────────────────────────────────────
// Helpers — consistent with failure_analysis.rs / missing_columns.rs
// ─────────────────────────────────────────────────────────────────────────────

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, Vec<String>> {
    map.iter()
        .map(|(k, v)| {
            let mut sorted = v.clone();
            sorted.sort();
            sorted.dedup();
            (k.clone(), sorted)
        })
        .collect()
}

fn normalize_to_short_names(map: &HashMap<String, Vec<String>>) -> HashMap<String, Vec<String>> {
    map.iter()
        .map(|(k, v)| {
            let normalized: Vec<String> = v
                .iter()
                .map(|ref_str| {
                    let parts: Vec<&str> =
                        ref_str.split('.').map(|p| p.trim_matches('"')).collect();
                    if parts.len() >= 3 {
                        let table = parts[parts.len() - 2];
                        let col = parts[parts.len() - 1];
                        format!("\"{table}\".\"{col}\"")
                    } else {
                        ref_str.clone()
                    }
                })
                .collect();
            let mut sorted = normalized;
            sorted.sort();
            sorted.dedup();
            (k.clone(), sorted)
        })
        .collect()
}

/// Collect all column names across every schema table (uppercased).
fn all_schema_columns(schema: &HashMap<String, serde_json::Value>) -> HashSet<String> {
    let mut set = HashSet::new();
    for val in schema.values() {
        if let serde_json::Value::Object(cols) = val {
            for col_name in cols.keys() {
                set.insert(col_name.to_uppercase());
            }
        }
    }
    set
}

/// Check whether SQL text contains SELECT * or alias.* patterns.
fn has_select_star(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    upper.contains("SELECT *")
        || upper.contains("SELECT\n*")
        || upper.contains("SELECT\r\n*")
        || upper.contains(".*")
}

/// Check whether SQL text contains a CTE (WITH ... AS).
fn has_cte(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    // Match "WITH name AS" but avoid false positives from TIMESTAMP WITH TIME ZONE etc.
    // Heuristic: WITH at the start of the query (possibly after whitespace) or after a closing paren.
    upper.trim_start().starts_with("WITH ")
        || upper.contains("\nWITH ")
        || upper.contains(") WITH ")
}

/// Check whether SQL text contains UNION or UNION ALL.
fn has_union(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    upper.contains("UNION ALL") || upper.contains("UNION\n") || upper.contains("UNION ")
}

/// Truncate SQL to `max` chars, replacing newlines with spaces.
fn sql_prefix(sql: &str, max: usize) -> String {
    let flat: String = sql
        .chars()
        .map(|c| if c == '\n' || c == '\r' { ' ' } else { c })
        .collect();
    if flat.len() > max {
        format!("{}...", &flat[..max])
    } else {
        flat
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Main
// ─────────────────────────────────────────────────────────────────────────────

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("failed to open fixture");
    let reader = std::io::BufReader::with_capacity(1024 * 1024, file);

    let dialect = SqlDialect::Snowflake;

    // ── Counters ─────────────────────────────────────────────────────────
    let mut total_cases = 0u64;
    let mut parse_fail = 0u64;
    let mut empty_actual = 0u64;
    let mut nonempty_with_missing = 0u64; // queries that qualify for this analysis

    let mut total_missing = 0u64;
    let mut cause_counts: HashMap<MissingCause, u64> = HashMap::new();

    // Per-query feature co-occurrence with missing columns
    let mut queries_with_star = 0u64;
    let mut queries_with_cte = 0u64;
    let mut queries_with_union = 0u64;

    // Per-column-name frequency of missing
    let mut missing_col_freq: HashMap<String, u64> = HashMap::new();

    // Samples per cause (up to 10 each)
    let mut cause_samples: HashMap<MissingCause, Vec<MissingEntry>> = HashMap::new();

    // ── Iterate cases ────────────────────────────────────────────────────
    for line in reader.lines() {
        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_cases += 1;

        if case.expected.is_empty() {
            continue;
        }

        let schema = schema_from_json(&case.schema);
        let schema_cols = all_schema_columns(&schema);

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

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

        // Skip if actual is empty — that is a different failure mode.
        if actual.is_empty() {
            empty_actual += 1;
            continue;
        }

        // Find expected columns NOT in actual.
        let missing: Vec<(&String, &Vec<String>)> = expected
            .iter()
            .filter(|(col, _)| !actual.contains_key(*col))
            .collect();

        if missing.is_empty() {
            continue;
        }

        // This query has a non-empty actual but is missing some expected columns.
        nonempty_with_missing += 1;

        let query_star = has_select_star(&case.sql);
        let query_cte = has_cte(&case.sql);
        let query_union = has_union(&case.sql);

        if query_star {
            queries_with_star += 1;
        }
        if query_cte {
            queries_with_cte += 1;
        }
        if query_union {
            queries_with_union += 1;
        }

        let prefix = sql_prefix(&case.sql, 100);
        let actual_keys: Vec<String> = actual.keys().cloned().collect();

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

            let col_upper = col.to_uppercase();
            let col_in_schema = schema_cols.contains(&col_upper);

            // Categorize
            let cause = if col_in_schema && query_star {
                MissingCause::StarExpansionFailure
            } else if col_in_schema {
                MissingCause::SchemaColumnMissed
            } else {
                MissingCause::ExpressionAlias
            };

            *cause_counts.entry(cause).or_default() += 1;

            let samples = cause_samples.entry(cause).or_default();
            if samples.len() < 10 {
                samples.push(MissingEntry {
                    query_name: case.name.clone(),
                    sql_prefix: prefix.clone(),
                    column: (*col).clone(),
                    expected_sources: (*sources).clone(),
                    actual_cols: actual_keys.clone(),
                    cause,
                    has_star: query_star,
                    has_cte: query_cte,
                    has_union: query_union,
                });
            }
        }
    }

    // ── Output ───────────────────────────────────────────────────────────
    let sep = "=".repeat(76);

    println!("\n{sep}");
    println!("  Missing Output Columns — non-empty actual, expected cols absent");
    println!("{sep}");

    // Overview
    println!("\n--- Overview ---");
    println!("Total test cases:                          {total_cases}");
    println!("Parse failures (skipped):                  {parse_fail}");
    println!("Empty actual (skipped, different bucket):  {empty_actual}");
    println!("Queries with non-empty actual + missing:   {nonempty_with_missing}");
    println!("Total missing output columns:              {total_missing}");
    if nonempty_with_missing > 0 {
        println!(
            "Avg missing cols per affected query:        {:.2}",
            total_missing as f64 / nonempty_with_missing as f64
        );
    }

    // Cause breakdown
    println!("\n--- Root Cause Breakdown ---");
    let mut causes: Vec<(MissingCause, u64)> = cause_counts.into_iter().collect();
    causes.sort_by(|a, b| b.1.cmp(&a.1));
    for (cause, count) in &causes {
        let pct = if total_missing > 0 {
            *count as f64 / total_missing as f64 * 100.0
        } else {
            0.0
        };
        println!("  {:45} {:>7}  ({:.1}%)", cause.label(), count, pct);
    }

    // SQL feature co-occurrence
    println!(
        "\n--- SQL Feature Co-occurrence (among {} affected queries) ---",
        nonempty_with_missing
    );
    if nonempty_with_missing > 0 {
        println!(
            "  Has SELECT * / alias.*:  {:>5}  ({:.1}%)",
            queries_with_star,
            queries_with_star as f64 / nonempty_with_missing as f64 * 100.0
        );
        println!(
            "  Has CTE (WITH ... AS):   {:>5}  ({:.1}%)",
            queries_with_cte,
            queries_with_cte as f64 / nonempty_with_missing as f64 * 100.0
        );
        println!(
            "  Has UNION / UNION ALL:   {:>5}  ({:.1}%)",
            queries_with_union,
            queries_with_union as f64 / nonempty_with_missing as f64 * 100.0
        );
    }

    // Top 20 most-frequently-missing column names
    println!("\n--- Top 20 Most-Frequently-Missing Column Names ---");
    let mut freq_vec: Vec<_> = missing_col_freq.into_iter().collect();
    freq_vec.sort_by(|a, b| b.1.cmp(&a.1));
    for (col, count) in freq_vec.iter().take(20) {
        println!("  {col:<45} {count:>6}");
    }

    // Samples per cause
    for (cause, _count) in &causes {
        let samples = match cause_samples.get(cause) {
            Some(s) if !s.is_empty() => s,
            _ => continue,
        };

        let n = samples.len().min(10);
        println!("\n{sep}\n  {} Sample Queries: {}\n{sep}", n, cause.label());

        for (i, entry) in samples.iter().take(10).enumerate() {
            println!("\n  [{:>2}] query:    {}", i + 1, entry.query_name);
            println!("       sql:      {}", entry.sql_prefix);
            println!("       missing:  {}", entry.column);
            println!("       expected: {:?}", entry.expected_sources);
            println!(
                "       actual:   [{}]",
                entry
                    .actual_cols
                    .iter()
                    .take(8)
                    .cloned()
                    .collect::<Vec<_>>()
                    .join(", ")
            );
            if entry.actual_cols.len() > 8 {
                println!(
                    "                 ... and {} more actual cols",
                    entry.actual_cols.len() - 8
                );
            }
            println!(
                "       features: star={} cte={} union={}",
                entry.has_star, entry.has_cte, entry.has_union
            );
        }
    }

    println!("\n{sep}");
    println!("  Done.");
    println!("{sep}\n");
}
