//! Analyze queries where we produce the CORRECT column set but WRONG sources.
//!
//! Filters to: `actual.keys() == expected.keys()` AND `actual != expected`.
//! For each such query, classifies every column into one of five categories:
//!   - exact:    actual sources == expected sources
//!   - subset:   actual is a strict subset of expected (missing sources)
//!   - superset: actual is a strict superset of expected (extra sources)
//!   - disjoint: actual and expected share NO overlap
//!   - overlap:  partial overlap (neither subset nor superset)
//!
//! Key question answered: how many queries have ALL columns as exact-or-subset?
//! Those are "almost right" — we never hallucinate extra sources, just miss some.
//!
//! Run: cargo run --example source_error_analysis --release

use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use serde::Deserialize;
use std::collections::{BTreeMap, 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>>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
enum ColumnCategory {
    Exact,
    Subset,
    Superset,
    Disjoint,
    Overlap,
}

impl std::fmt::Display for ColumnCategory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ColumnCategory::Exact => write!(f, "exact"),
            ColumnCategory::Subset => write!(f, "subset"),
            ColumnCategory::Superset => write!(f, "superset"),
            ColumnCategory::Disjoint => write!(f, "disjoint"),
            ColumnCategory::Overlap => write!(f, "overlap"),
        }
    }
}

/// Per-column detail for a subset column in a sample query.
struct SubsetColumnDetail {
    column_name: String,
    actual_sources: Vec<String>,
    missing_sources: Vec<String>,
}

/// A sample "subset-only" query for detailed display.
struct SubsetOnlySample {
    name: String,
    sql_prefix: String,
    details: Vec<SubsetColumnDetail>,
}

// ===========================================================================
// 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 each source list; leave keys as-is.
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()
}

/// Trim each source ref to the last two dot-separated parts: `"TABLE"."COL"`.
/// Then sort + dedup.
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 the table portion from a `"TABLE"."COL"` reference (uppercased).
fn table_of(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 {
        ref_str.trim_matches('"').to_uppercase()
    }
}

/// Classify a single column given its actual and expected source sets.
fn classify_column(actual: &[String], expected: &[String]) -> ColumnCategory {
    let act_set: HashSet<&String> = actual.iter().collect();
    let exp_set: HashSet<&String> = expected.iter().collect();

    if act_set == exp_set {
        return ColumnCategory::Exact;
    }

    let intersection_count = act_set.intersection(&exp_set).count();

    if intersection_count == 0 {
        // Check edge case: both empty would be exact (handled above).
        return ColumnCategory::Disjoint;
    }

    let act_subset_of_exp = act_set.is_subset(&exp_set);
    let exp_subset_of_act = exp_set.is_subset(&act_set);

    if act_subset_of_exp {
        ColumnCategory::Subset
    } else if exp_subset_of_act {
        ColumnCategory::Superset
    } else {
        ColumnCategory::Overlap
    }
}

// ===========================================================================
// 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_queries = 0u64;
    let mut parse_fail = 0u64;
    let mut exact_queries = 0u64;

    // Queries where actual.keys() == expected.keys() but values differ.
    let mut same_keys_diff_sources = 0u64;

    // Per-column category counts (across all same-keys-diff-sources queries).
    let mut col_category_counts: BTreeMap<ColumnCategory, u64> = BTreeMap::new();
    let mut total_columns_analyzed = 0u64;

    // Queries where ALL columns are either exact or subset (no extras ever).
    let mut subset_only_queries = 0u64;

    // For subset-only queries: per-column stats.
    let mut subset_only_exact_cols = 0u64;
    let mut subset_only_subset_cols = 0u64;

    // Distribution of missing-source-count per subset column (across subset-only queries).
    let mut missing_count_dist: BTreeMap<usize, u64> = BTreeMap::new();

    // Distribution of missing source TABLE name (across subset-only queries).
    let mut missing_table_freq: HashMap<String, u64> = HashMap::new();

    // Collect first 20 subset-only samples for detailed printing.
    let mut subset_only_samples: Vec<SubsetOnlySample> = Vec::new();

    // -----------------------------------------------------------------------
    // Process
    // -----------------------------------------------------------------------

    for line in reader.lines() {
        let line = line.expect("failed to read line");
        let line = line.trim();
        if line.is_empty() {
            continue;
        }

        let case: Case = match serde_json::from_str(line) {
            Ok(c) => c,
            Err(_) => continue,
        };
        total_queries += 1;

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

        // Check: exact match?
        if actual == expected {
            exact_queries += 1;
            continue;
        }

        // Check: same key set?
        let actual_keys: HashSet<&String> = actual.keys().collect();
        let expected_keys: HashSet<&String> = expected.keys().collect();
        if actual_keys != expected_keys {
            continue;
        }

        // Same keys, different sources.
        same_keys_diff_sources += 1;

        // Classify each column.
        let mut all_exact_or_subset = true;
        let mut query_details: Vec<SubsetColumnDetail> = Vec::new();

        for (col, exp_sources) in &expected {
            let act_sources = actual.get(col).unwrap();
            let category = classify_column(act_sources, exp_sources);

            *col_category_counts.entry(category).or_default() += 1;
            total_columns_analyzed += 1;

            if category != ColumnCategory::Exact && category != ColumnCategory::Subset {
                all_exact_or_subset = false;
            }
        }

        if !all_exact_or_subset {
            continue;
        }

        // This is a subset-only query.
        subset_only_queries += 1;

        // Detailed stats for subset-only queries.
        for (col, exp_sources) in &expected {
            let act_sources = actual.get(col).unwrap();
            let category = classify_column(act_sources, exp_sources);

            match category {
                ColumnCategory::Exact => {
                    subset_only_exact_cols += 1;
                }
                ColumnCategory::Subset => {
                    subset_only_subset_cols += 1;

                    let act_set: HashSet<&String> = act_sources.iter().collect();
                    let missing: Vec<String> = exp_sources
                        .iter()
                        .filter(|e| !act_set.contains(e))
                        .cloned()
                        .collect();

                    *missing_count_dist.entry(missing.len()).or_default() += 1;

                    for miss_ref in &missing {
                        let tbl = table_of(miss_ref);
                        if !tbl.is_empty() {
                            *missing_table_freq.entry(tbl).or_default() += 1;
                        }
                    }

                    // Collect detail for samples.
                    if subset_only_samples.len() < 20
                        || subset_only_samples
                            .last()
                            .map_or(true, |s| s.name == case.name)
                    {
                        query_details.push(SubsetColumnDetail {
                            column_name: col.clone(),
                            actual_sources: act_sources.clone(),
                            missing_sources: missing,
                        });
                    }
                }
                _ => {} // unreachable for subset-only queries
            }
        }

        if subset_only_samples.len() < 20 && !query_details.is_empty() {
            let sql_prefix: String = case
                .sql
                .chars()
                .filter(|c| !c.is_ascii_control() || *c == ' ')
                .take(200)
                .collect();

            subset_only_samples.push(SubsetOnlySample {
                name: case.name.clone(),
                sql_prefix,
                details: query_details,
            });
        }
    }

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

    println!();
    println!("{}", "=".repeat(78));
    println!("  SOURCE ERROR ANALYSIS");
    println!("  Queries with correct column set but wrong sources");
    println!("{}", "=".repeat(78));

    println!();
    println!("--- Query-Level Overview ---");
    println!("Total queries:                   {total_queries}");
    println!("Parse failures:                  {parse_fail}");
    println!("Exact match:                     {exact_queries}");
    println!(
        "Same keys, different sources:    {same_keys_diff_sources} ({:.1}% of non-exact parsed)",
        same_keys_diff_sources as f64 / (total_queries - parse_fail - exact_queries).max(1) as f64
            * 100.0
    );

    println!();
    println!(
        "--- Per-Column Category Distribution ({total_columns_analyzed} columns across {same_keys_diff_sources} queries) ---"
    );
    let categories = [
        ColumnCategory::Exact,
        ColumnCategory::Subset,
        ColumnCategory::Superset,
        ColumnCategory::Disjoint,
        ColumnCategory::Overlap,
    ];
    for cat in &categories {
        let count = col_category_counts.get(cat).copied().unwrap_or(0);
        println!(
            "  {:<12} {:>7} ({:5.1}%)",
            cat.to_string(),
            count,
            count as f64 / total_columns_analyzed.max(1) as f64 * 100.0
        );
    }

    println!();
    println!("--- Subset-Only Queries (all columns exact or subset, never extra) ---");
    println!(
        "Count:          {subset_only_queries} ({:.1}% of same-keys queries)",
        subset_only_queries as f64 / same_keys_diff_sources.max(1) as f64 * 100.0
    );
    println!("  Exact cols:   {subset_only_exact_cols}");
    println!("  Subset cols:  {subset_only_subset_cols}");

    println!();
    println!("--- Missing Source Count Distribution (per subset column, subset-only queries) ---");
    for (n, count) in &missing_count_dist {
        println!("  {n} missing:  {count} columns");
    }

    println!();
    println!("--- Top 20 Missing Source TABLEs (subset-only queries) ---");
    let mut tbl_vec: Vec<_> = missing_table_freq.into_iter().collect();
    tbl_vec.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
    for (i, (tbl, count)) in tbl_vec.iter().take(20).enumerate() {
        println!("  {:>2}. {tbl:<45} {count:>6}", i + 1);
    }

    // Print first 20 subset-only sample queries.
    println!();
    println!("{}", "=".repeat(78));
    println!("  FIRST 20 SUBSET-ONLY QUERIES (all columns exact or subset)");
    println!("{}", "=".repeat(78));

    for (i, sample) in subset_only_samples.iter().enumerate() {
        println!();
        println!("--- #{} ---", i + 1);
        println!("Name: {}", sample.name);
        println!("SQL:  {}...", sample.sql_prefix);

        for detail in &sample.details {
            println!("  Column: {}", detail.column_name);
            println!(
                "    Actual sources:  [{}]",
                detail.actual_sources.join(", ")
            );
            println!(
                "    Missing sources: [{}]",
                detail.missing_sources.join(", ")
            );
        }
    }

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