//! Deep analysis of the SUBSET category — columns where actual lineage is a strict
//! subset of expected (we have the output column, but are missing source references).
//!
//! This diagnostic answers:
//!   - Are missing sources from the SAME table as existing actual sources, or DIFFERENT?
//!   - Do missing source columns exist in the query's schema at all?
//!   - What are the top missing column names?
//!   - Which SQL constructs correlate with subset mismatches?
//!   - How many queries would become exact if subset gaps were closed?
//!   - Which queries have the widest gap (most missing refs)?
//!
//! Run: cargo run --example subset_analysis --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;

// ===========================================================================
// Fixture type
// ===========================================================================

#[derive(Debug, Deserialize)]
struct Case {
    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(),
    }
}

/// Normalize a column-lineage map: sort + dedup each source list.
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()
}

/// Normalize actual output to short `"TABLE"."COL"` form (strip leading db/schema qualifiers).
fn normalize_to_short(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 {
        String::new()
    }
}

/// Extract the column portion from a `"TABLE"."COL"` reference (uppercased).
fn column_of(ref_str: &str) -> String {
    let parts: Vec<&str> = ref_str.split('.').map(|p| p.trim_matches('"')).collect();
    parts.last().unwrap_or(&"").to_uppercase()
}

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

/// Case-insensitive check for whether SQL contains a keyword.
/// Uses simple substring matching on uppercased SQL.
fn sql_has(sql_upper: &str, keyword: &str) -> bool {
    sql_upper.contains(keyword)
}

// ===========================================================================
// Per-query tracking for gap analysis
// ===========================================================================

/// Tracks subset mismatch details for a single query.
struct QuerySubsetInfo {
    name: String,
    sql: String,
    /// Total missing source refs across all subset columns in this query.
    total_missing: usize,
    /// Number of columns that are subset (have output col, missing some sources).
    subset_column_count: usize,
    /// Whether ALL per-column mismatches in this query are subset-type
    /// (no superset, no extra output, no missing output — only subset + exact).
    all_mismatches_are_subset: bool,
    /// The column with the most missing refs and its expected/actual.
    worst_column: String,
    worst_expected: Vec<String>,
    worst_actual: Vec<String>,
}

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

    // Subset-specific counters (per missing source ref)
    let mut total_missing_refs = 0u64;
    let mut missing_same_table = 0u64; // Missing source from a table we DO have in actual
    let mut missing_diff_table = 0u64; // Missing source from a table NOT in any actual source
    let mut missing_col_in_schema = 0u64; // Missing source's column exists somewhere in schema
    let mut missing_col_not_in_schema = 0u64; // Missing source's column not in schema at all

    // SQL construct correlation (counted per query that has at least one subset column)
    let mut queries_with_subset = 0u64;
    let mut has_cte = 0u64;
    let mut has_subquery = 0u64;
    let mut has_union = 0u64;
    let mut has_join = 0u64;
    let mut has_case = 0u64;
    let mut has_window = 0u64;

    // Top missing column names (case-insensitive)
    let mut missing_col_freq: HashMap<String, u64> = HashMap::new();

    // Per-query gap tracking for worst-case samples and exact-if-fixed count
    let mut query_infos: Vec<QuerySubsetInfo> = Vec::new();

    // -----------------------------------------------------------------------
    // Process each fixture
    // -----------------------------------------------------------------------

    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 schema_cols = all_schema_columns(&schema);

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

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

        if actual == expected {
            exact_queries += 1;
            continue;
        }

        // Per-column classification for this query
        let mut query_subset_count = 0usize;
        let mut query_missing_total = 0usize;
        let mut has_non_subset_mismatch = false;

        // Track the worst column (most missing refs) in this query
        let mut worst_col = String::new();
        let mut worst_expected: Vec<String> = Vec::new();
        let mut worst_actual: Vec<String> = Vec::new();
        let mut worst_missing_count = 0usize;

        // Check expected columns
        for (col, exp_sources) in &expected {
            match actual.get(col) {
                None => {
                    // Missing output column entirely — not a subset case
                    has_non_subset_mismatch = true;
                }
                Some(act_sources) => {
                    let all_actual_in_expected =
                        act_sources.iter().all(|a| exp_sources.contains(a));
                    let all_expected_in_actual =
                        exp_sources.iter().all(|e| act_sources.contains(e));

                    if all_expected_in_actual && all_actual_in_expected {
                        // Exact match for this column
                        continue;
                    }

                    if !all_actual_in_expected {
                        // Actual has sources NOT in expected (superset or mixed)
                        has_non_subset_mismatch = true;
                        // If also missing some expected, still count the missing part
                        // but flag that the query isn't purely subset
                    }

                    if all_expected_in_actual {
                        // All expected present — this is superset, not subset
                        has_non_subset_mismatch = true;
                        continue;
                    }

                    // At least some expected sources are missing from actual.
                    // This column contributes to the subset category.
                    let missing: Vec<&String> = exp_sources
                        .iter()
                        .filter(|e| !act_sources.contains(e))
                        .collect();

                    if missing.is_empty() {
                        continue;
                    }

                    // Only count as pure subset if actual is a strict subset
                    // (all actual sources are in expected, but not all expected in actual).
                    if !all_actual_in_expected {
                        // Mixed: has extra AND missing. Still analyze the missing part.
                        has_non_subset_mismatch = true;
                    }

                    query_subset_count += 1;
                    query_missing_total += missing.len();

                    // Tables present in actual sources for this column
                    let actual_tables: HashSet<String> =
                        act_sources.iter().map(|s| table_of(s)).collect();

                    for miss_ref in &missing {
                        total_missing_refs += 1;

                        let miss_table = table_of(miss_ref);
                        let miss_col = column_of(miss_ref);

                        // Same table vs different table
                        if actual_tables.contains(&miss_table) {
                            missing_same_table += 1;
                        } else {
                            missing_diff_table += 1;
                        }

                        // Column exists in schema?
                        if schema_cols.contains(&miss_col) {
                            missing_col_in_schema += 1;
                        } else {
                            missing_col_not_in_schema += 1;
                        }

                        // Frequency of missing column name
                        *missing_col_freq.entry(miss_col).or_default() += 1;
                    }

                    // Track worst column for this query
                    if missing.len() > worst_missing_count {
                        worst_missing_count = missing.len();
                        worst_col = col.clone();
                        worst_expected = exp_sources.clone();
                        worst_actual = act_sources.clone();
                    }
                }
            }
        }

        // Check for extra output columns (actual has columns not in expected)
        for col in actual.keys() {
            if !expected.contains_key(col) {
                has_non_subset_mismatch = true;
            }
        }

        // SQL construct correlation (only for queries with subset columns)
        if query_subset_count > 0 {
            queries_with_subset += 1;

            let sql_upper = case.sql.to_uppercase();
            if sql_has(&sql_upper, "WITH ") {
                has_cte += 1;
            }
            // Subquery: opening paren followed by SELECT (rough heuristic)
            if sql_upper.contains("(SELECT ") || sql_upper.contains("( SELECT ") {
                has_subquery += 1;
            }
            if sql_has(&sql_upper, "UNION") {
                has_union += 1;
            }
            if sql_has(&sql_upper, " JOIN ") {
                has_join += 1;
            }
            if sql_has(&sql_upper, "CASE ") {
                has_case += 1;
            }
            if sql_has(&sql_upper, " OVER(") || sql_has(&sql_upper, " OVER (") {
                has_window += 1;
            }

            query_infos.push(QuerySubsetInfo {
                name: case.name.clone(),
                sql: case.sql.clone(),
                total_missing: query_missing_total,
                subset_column_count: query_subset_count,
                all_mismatches_are_subset: !has_non_subset_mismatch,
                worst_column: worst_col,
                worst_expected,
                worst_actual,
            });
        }
    }

    // -----------------------------------------------------------------------
    // Derived metrics
    // -----------------------------------------------------------------------

    let would_become_exact = query_infos
        .iter()
        .filter(|q| q.all_mismatches_are_subset)
        .count() as u64;

    // Sort by total_missing descending for worst-case samples
    query_infos.sort_by(|a, b| b.total_missing.cmp(&a.total_missing));

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

    println!();
    println!("{}", "=".repeat(74));
    println!("  SUBSET ANALYSIS — Missing Source References in Column Lineage");
    println!("{}", "=".repeat(74));

    println!();
    println!("--- Query-Level Overview ---");
    println!("Total queries:             {total_queries}");
    println!("Exact match:               {exact_queries}");
    println!("Parse failures:            {parse_fail_queries}");
    println!(
        "Queries with subset cols:  {queries_with_subset} ({:.1}%)",
        queries_with_subset as f64 / total_queries as f64 * 100.0
    );

    println!();
    println!("--- Missing Source Reference Totals ---");
    println!("Total missing source refs: {total_missing_refs}");
    if queries_with_subset > 0 {
        println!(
            "Avg missing per query:     {:.1}",
            total_missing_refs as f64 / queries_with_subset as f64
        );
    }

    println!();
    println!("--- Where Are Missing Sources? (per missing ref) ---");
    let pct = |n: u64| -> f64 {
        if total_missing_refs > 0 {
            n as f64 / total_missing_refs as f64 * 100.0
        } else {
            0.0
        }
    };
    println!(
        "  Same table as actual:      {missing_same_table:>7}  ({:5.1}%)",
        pct(missing_same_table)
    );
    println!(
        "  Different table:           {missing_diff_table:>7}  ({:5.1}%)",
        pct(missing_diff_table)
    );

    println!();
    println!("--- Missing Column in Schema? (per missing ref) ---");
    println!(
        "  Column exists in schema:   {missing_col_in_schema:>7}  ({:5.1}%)",
        pct(missing_col_in_schema)
    );
    println!(
        "  Column NOT in schema:      {missing_col_not_in_schema:>7}  ({:5.1}%)",
        pct(missing_col_not_in_schema)
    );

    // SQL construct correlation
    println!();
    println!("--- SQL Construct Correlation (among {queries_with_subset} subset queries) ---");
    let qpct = |n: u64| -> f64 {
        if queries_with_subset > 0 {
            n as f64 / queries_with_subset as f64 * 100.0
        } else {
            0.0
        }
    };
    println!(
        "  Has CTE (WITH):         {has_cte:>7}  ({:5.1}%)",
        qpct(has_cte)
    );
    println!(
        "  Has subquery:           {has_subquery:>7}  ({:5.1}%)",
        qpct(has_subquery)
    );
    println!(
        "  Has UNION:              {has_union:>7}  ({:5.1}%)",
        qpct(has_union)
    );
    println!(
        "  Has JOIN:               {has_join:>7}  ({:5.1}%)",
        qpct(has_join)
    );
    println!(
        "  Has CASE:               {has_case:>7}  ({:5.1}%)",
        qpct(has_case)
    );
    println!(
        "  Has window fn (OVER):   {has_window:>7}  ({:5.1}%)",
        qpct(has_window)
    );

    // Impact: queries that would become exact
    println!();
    println!("--- Impact ---");
    println!(
        "Queries that become exact if subset fixed: {would_become_exact} ({:.1}% of total)",
        would_become_exact as f64 / total_queries as f64 * 100.0
    );
    println!(
        "  (These queries have ONLY subset mismatches — no extra cols, no superset, no missing output)"
    );

    // Top 20 missing column names
    println!();
    println!("--- Top 20 Most Frequently Missing Source Column Names ---");
    let mut col_vec: Vec<_> = missing_col_freq.into_iter().collect();
    col_vec.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
    for (i, (col, count)) in col_vec.iter().take(20).enumerate() {
        println!("  {:>2}. {col:<40} {count:>6}", i + 1);
    }

    // 5 worst-case sample queries
    println!();
    println!("{}", "=".repeat(74));
    println!("  5 QUERIES WITH LARGEST SUBSET GAP (most missing source refs)");
    println!("{}", "=".repeat(74));

    for (i, info) in query_infos.iter().take(5).enumerate() {
        let sql_prefix: String = info
            .sql
            .chars()
            .filter(|c| !c.is_ascii_control() || *c == ' ')
            .take(200)
            .collect();

        println!();
        println!("--- #{} ---", i + 1);
        println!("Query:           {}", info.name);
        println!(
            "Missing refs:    {} across {} subset columns",
            info.total_missing, info.subset_column_count
        );
        println!(
            "Pure subset?     {}",
            if info.all_mismatches_are_subset {
                "yes (would become exact)"
            } else {
                "no (has other mismatch types too)"
            }
        );
        println!("SQL:             {sql_prefix}...");
        println!();
        println!("  Worst column: {}", info.worst_column);
        println!("    Expected: [{}]", info.worst_expected.join(", "));
        println!("    Actual:   [{}]", info.worst_actual.join(", "));

        // Show the missing refs for the worst column
        let missing: Vec<&String> = info
            .worst_expected
            .iter()
            .filter(|e| !info.worst_actual.contains(e))
            .collect();
        println!(
            "    Missing:  [{}]",
            missing
                .iter()
                .map(|s| s.as_str())
                .collect::<Vec<_>>()
                .join(", ")
        );
    }

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