//! Analyze the 59,920 expected columns missing from actual lineage output.
//!
//! For each test case, compares expected vs actual column lineage and
//! categorizes every missing column by root cause:
//!   - SELECT * not expanded (column exists in schema but no edge emitted)
//!   - Expression alias (output column is a computed alias, not a schema column)
//!   - Unknown (everything else)
//!
//! Run: cargo run --example missing_columns --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>>,
}

#[derive(Debug, Clone)]
struct MissingSample {
    sql_prefix: String,
    column: String,
    expected_sources: Vec<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum RootCause {
    /// Column exists in schema, query uses SELECT * or alias.*, but no edge emitted.
    SelectStarNotExpanded,
    /// Column exists in schema, query does NOT use SELECT *, still no edge.
    SchemaColumnNoEdge,
    /// Column is an expression alias (not in any schema table).
    ExpressionAlias,
    /// Expected sources reference tables not in schema.
    SourceTableNotInSchema,
    /// Parse or lineage failure (empty actual output).
    ParseFailure,
    /// Other / unknown.
    Unknown,
}

impl RootCause {
    fn label(&self) -> &'static str {
        match self {
            RootCause::SelectStarNotExpanded => "SELECT * not expanded",
            RootCause::SchemaColumnNoEdge => "Schema column, no edge (non-star)",
            RootCause::ExpressionAlias => "Expression alias",
            RootCause::SourceTableNotInSchema => "Source table not in schema",
            RootCause::ParseFailure => "Parse / lineage failure",
            RootCause::Unknown => "Unknown",
        }
    }
}

// ───────────────────────────────────────────────────────────────────────────────
// Helpers — same patterns as failure_analysis.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()
}

/// Extract table name from `"TABLE"."COL"`.
fn table_only(ref_str: &str) -> String {
    let parts: Vec<&str> = ref_str.split('.').map(|p| p.trim_matches('"')).collect();
    if parts.len() >= 2 {
        parts[parts.len() - 2].to_uppercase()
    } else {
        String::new()
    }
}

/// Check whether SQL text contains any SELECT * or alias.* pattern.
fn has_select_star(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    // Match SELECT *, SELECT t.*, SELECT alias.*
    upper.contains("SELECT *") || upper.contains(".*")
}

/// Build a set of all column names across all schema tables (uppercased).
fn all_schema_columns(
    schema_columns: &HashMap<String, Vec<String>>,
) -> std::collections::HashSet<String> {
    schema_columns
        .values()
        .flat_map(|cols| cols.iter().cloned())
        .collect()
}

/// Check if an expected column's source tables all exist in the schema.
fn sources_in_schema(sources: &[String], schema_tables: &[String]) -> bool {
    sources.iter().all(|src| {
        let tbl = table_only(src);
        tbl.is_empty() || schema_tables.iter().any(|s| s.eq_ignore_ascii_case(&tbl))
    })
}

// ───────────────────────────────────────────────────────────────────────────────
// 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;

    // Accumulators
    let mut total_cases = 0u64;
    let mut total_expected_cols = 0u64;
    let mut total_actual_cols = 0u64;
    let mut total_missing = 0u64;

    // Per root cause: count + sample cases
    let mut cause_counts: HashMap<RootCause, u64> = HashMap::new();
    let mut cause_samples: HashMap<RootCause, Vec<MissingSample>> = HashMap::new();

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

    // Wildcard reference tracker (expected sources containing "*")
    let mut wildcard_count = 0u64;

    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;

        let schema = schema_from_json(&case.schema);
        let schema_tables: Vec<String> = schema.keys().map(|k| k.to_uppercase()).collect();

        // Build column map per table
        let mut schema_columns: HashMap<String, Vec<String>> = HashMap::new();
        for (tbl, val) in &schema {
            if let serde_json::Value::Object(cols) = val {
                schema_columns.insert(
                    tbl.to_uppercase(),
                    cols.keys().map(|c| c.to_uppercase()).collect(),
                );
            }
        }

        let all_cols = all_schema_columns(&schema_columns);
        let expected = normalize(&case.expected);
        total_expected_cols += expected.len() as u64;

        let query_has_star = has_select_star(&case.sql);
        let sql_prefix = if case.sql.len() > 120 {
            format!("{}...", &case.sql[..120])
        } else {
            case.sql.clone()
        }
        .replace('\n', " ");

        // Run lineage
        let actual = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => normalize_to_short_names(&r.column_dict),
            Err(_) => {
                // Parse failure: all expected columns are missing
                for (col, sources) in &expected {
                    total_missing += 1;
                    *missing_col_freq.entry(col.clone()).or_default() += 1;

                    // Check wildcard
                    if sources.iter().any(|s| s.contains('*')) {
                        wildcard_count += 1;
                    }

                    let cause = RootCause::ParseFailure;
                    *cause_counts.entry(cause).or_default() += 1;
                    let samples = cause_samples.entry(cause).or_default();
                    if samples.len() < 10 {
                        samples.push(MissingSample {
                            sql_prefix: sql_prefix.clone(),
                            column: col.clone(),
                            expected_sources: sources.clone(),
                        });
                    }
                }
                continue;
            }
        };

        total_actual_cols += actual.len() as u64;

        // Find expected columns NOT in actual
        for (col, sources) in &expected {
            if actual.contains_key(col) {
                continue;
            }

            // This column is missing from actual output
            total_missing += 1;
            *missing_col_freq.entry(col.clone()).or_default() += 1;

            // Check wildcard in expected sources
            if sources.iter().any(|s| s.contains('*')) {
                wildcard_count += 1;
            }

            // Determine root cause
            let col_upper = col.to_uppercase();
            let col_in_schema = all_cols.contains(&col_upper);
            let srcs_in_schema = sources_in_schema(sources, &schema_tables);

            let cause = if !srcs_in_schema {
                RootCause::SourceTableNotInSchema
            } else if col_in_schema && query_has_star {
                RootCause::SelectStarNotExpanded
            } else if col_in_schema && !query_has_star {
                RootCause::SchemaColumnNoEdge
            } else if !col_in_schema {
                RootCause::ExpressionAlias
            } else {
                RootCause::Unknown
            };

            *cause_counts.entry(cause).or_default() += 1;
            let samples = cause_samples.entry(cause).or_default();
            if samples.len() < 10 {
                samples.push(MissingSample {
                    sql_prefix: sql_prefix.clone(),
                    column: col.clone(),
                    expected_sources: sources.clone(),
                });
            }
        }
    }

    // ───────────────────────────────────────────────────────────────────────
    // Output
    // ───────────────────────────────────────────────────────────────────────

    let sep = "=".repeat(72);
    println!("\n{sep}");
    println!("  Missing Column Analysis — Expected in ground truth, absent in actual");
    println!("{sep}");

    println!("\n--- Overview ---");
    println!("Total test cases:        {total_cases}");
    println!("Total expected columns:  {total_expected_cols}");
    println!("Total actual columns:    {total_actual_cols}");
    println!("Total MISSING columns:   {total_missing}");
    println!(
        "Missing rate:            {:.1}%",
        total_missing as f64 / total_expected_cols as f64 * 100.0
    );
    println!("Wildcard refs in expected sources: {wildcard_count}");

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

    // Top missing column names
    println!("\n--- Top 30 Missing Column Names (by frequency) ---");
    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(30) {
        println!("  {col:40} {count:>6}");
    }

    // Samples per root cause
    for (cause, _count) in &causes {
        let samples = match cause_samples.get(cause) {
            Some(s) => s,
            None => continue,
        };
        if samples.is_empty() {
            continue;
        }
        println!("\n--- 10 Samples: {} ---", cause.label());
        for (i, s) in samples.iter().enumerate() {
            println!(
                "  [{:>2}] column: {}\n       sources: {:?}\n       sql: {}",
                i + 1,
                s.column,
                s.expected_sources,
                s.sql_prefix
            );
        }
    }

    println!();
}
