//! Validate the STORE_ID extra output column cases.
//!
//! For each case where our engine produces STORE_ID but the ground truth doesn't,
//! determine: Is STORE_ID in the schema? Is it explicitly referenced in SQL?
//! Is it coming from SELECT * expansion? What does the SQL look like?
//!
//! Run: `cargo run --release --example validate_store_id`

use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use std::collections::{BTreeSet, HashMap};
use std::io::BufRead;

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

#[derive(serde::Serialize)]
struct StoreIdCase {
    name: String,
    sql_snippet: String,
    has_select_star: bool,
    has_explicit_store_id: bool,
    store_id_in_schema_tables: Vec<String>,
    expected_cols: Vec<String>,
    actual_cols: Vec<String>,
    extra_cols: Vec<String>,
    store_id_sources: Vec<String>,
    other_extra_cols: Vec<String>,
    would_match_without_extras: bool,
}

fn schema_tables_with_store_id(schema: &serde_json::Value) -> Vec<String> {
    let mut tables = Vec::new();
    if let Some(obj) = schema.as_object() {
        for (k, v) in obj {
            if let Some(cols) = v.as_object() {
                if cols.values().all(|vv| vv.is_string()) {
                    // flat table
                    if cols.keys().any(|c| c.to_uppercase() == "STORE_ID") {
                        tables.push(k.to_uppercase());
                    }
                } else {
                    // nested schema.table
                    for (k2, v2) in cols {
                        if let Some(inner) = v2.as_object() {
                            if inner.values().all(|vv| vv.is_string()) {
                                if inner.keys().any(|c| c.to_uppercase() == "STORE_ID") {
                                    tables.push(format!(
                                        "{}.{}",
                                        k.to_uppercase(),
                                        k2.to_uppercase()
                                    ));
                                }
                            } else {
                                for (k3, v3) in inner {
                                    if let Some(deepest) = v3.as_object() {
                                        if deepest.keys().any(|c| c.to_uppercase() == "STORE_ID") {
                                            tables.push(format!(
                                                "{}.{}.{}",
                                                k.to_uppercase(),
                                                k2.to_uppercase(),
                                                k3.to_uppercase()
                                            ));
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    tables.sort();
    tables
}

fn to_schema_map(val: &serde_json::Value) -> HashMap<String, serde_json::Value> {
    let mut m = HashMap::new();
    if let Some(obj) = val.as_object() {
        for (k, v) in obj {
            m.insert(k.clone(), v.clone());
        }
    }
    m
}

fn main() {
    let fixture_path = "tests/fixtures/lineage/snowflake_ground_truth.jsonl";
    let file = std::fs::File::open(fixture_path).expect("open fixture");
    let reader = std::io::BufReader::new(file);

    let mut cases: Vec<StoreIdCase> = Vec::new();
    let mut total = 0;

    for line in reader.lines() {
        let line = line.unwrap();
        if line.trim().is_empty() {
            continue;
        }
        total += 1;

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

        let schema_map = to_schema_map(&case.schema);
        let result = match column_lineage(&case.sql, SqlDialect::Snowflake, &schema_map, None) {
            Ok(r) => r,
            Err(_) => continue,
        };

        let expected_cols: BTreeSet<String> =
            case.expected.keys().map(|k| k.to_uppercase()).collect();
        let actual_cols: BTreeSet<String> = result
            .column_dict
            .keys()
            .map(|k| k.to_uppercase())
            .collect();
        let extra_cols: BTreeSet<String> =
            actual_cols.difference(&expected_cols).cloned().collect();

        if !extra_cols.contains("STORE_ID") {
            continue;
        }

        let sql_upper = case.sql.to_uppercase();
        let has_select_star = sql_upper.contains("SELECT *")
            || sql_upper.contains("SELECT\n*")
            || sql_upper.contains("SELECT\r\n*");
        // Check for explicit STORE_ID reference (not just in SELECT *)
        let has_explicit = sql_upper.contains("STORE_ID");

        let store_id_tables = schema_tables_with_store_id(&case.schema);

        let store_id_sources: Vec<String> = result
            .column_dict
            .iter()
            .find(|(k, _)| k.to_uppercase() == "STORE_ID")
            .map(|(_, v)| v.iter().cloned().collect())
            .unwrap_or_default();

        let other_extras: Vec<String> = extra_cols
            .iter()
            .filter(|c| *c != "STORE_ID")
            .cloned()
            .collect();

        // Would this case be exact match if we removed all extra cols?
        let missing_cols: BTreeSet<String> =
            expected_cols.difference(&actual_cols).cloned().collect();
        let mut would_match = missing_cols.is_empty();
        if would_match {
            // Check sources too
            for (col, expected_srcs) in &case.expected {
                let col_upper = col.to_uppercase();
                if let Some(actual_srcs) = result
                    .column_dict
                    .iter()
                    .find(|(k, _)| k.to_uppercase() == col_upper)
                    .map(|(_, v)| v)
                {
                    let exp_set: BTreeSet<String> =
                        expected_srcs.iter().map(|s| s.to_uppercase()).collect();
                    let act_set: BTreeSet<String> =
                        actual_srcs.iter().map(|s| s.to_uppercase()).collect();
                    if exp_set != act_set {
                        would_match = false;
                        break;
                    }
                }
            }
        }

        let sql_snippet = if case.sql.len() > 300 {
            format!("{}...", &case.sql[..300])
        } else {
            case.sql.clone()
        };

        cases.push(StoreIdCase {
            name: case.name,
            sql_snippet,
            has_select_star: has_select_star || sql_upper.contains(".*"),
            has_explicit_store_id: has_explicit,
            store_id_in_schema_tables: store_id_tables,
            expected_cols: expected_cols.into_iter().collect(),
            actual_cols: actual_cols.into_iter().collect(),
            extra_cols: extra_cols.into_iter().collect(),
            store_id_sources,
            other_extra_cols: other_extras,
            would_match_without_extras: would_match,
        });
    }

    eprintln!("Total queries: {}", total);
    eprintln!("STORE_ID extra cases: {}", cases.len());

    // Categorize
    let with_star = cases.iter().filter(|c| c.has_select_star).count();
    let with_explicit = cases.iter().filter(|c| c.has_explicit_store_id).count();
    let would_match = cases
        .iter()
        .filter(|c| c.would_match_without_extras)
        .count();
    let only_store_id_extra = cases
        .iter()
        .filter(|c| c.other_extra_cols.is_empty())
        .count();

    eprintln!("\n=== Summary ===");
    eprintln!(
        "  Has SELECT * or alias.*: {} ({:.1}%)",
        with_star,
        with_star as f64 / cases.len() as f64 * 100.0
    );
    eprintln!(
        "  Has explicit STORE_ID in SQL: {} ({:.1}%)",
        with_explicit,
        with_explicit as f64 / cases.len() as f64 * 100.0
    );
    eprintln!(
        "  Would be exact match if extras removed: {} ({:.1}%)",
        would_match,
        would_match as f64 / cases.len() as f64 * 100.0
    );
    eprintln!(
        "  Only extra col is STORE_ID (no others): {} ({:.1}%)",
        only_store_id_extra,
        only_store_id_extra as f64 / cases.len() as f64 * 100.0
    );

    // How many schema tables have STORE_ID
    let mut table_freq: HashMap<String, usize> = HashMap::new();
    for c in &cases {
        for t in &c.store_id_in_schema_tables {
            *table_freq.entry(t.clone()).or_default() += 1;
        }
    }
    let mut table_freq: Vec<_> = table_freq.into_iter().collect();
    table_freq.sort_by(|a, b| b.1.cmp(&a.1));

    eprintln!("\n=== Schema tables with STORE_ID column (top 20) ===");
    for (t, cnt) in table_freq.iter().take(20) {
        eprintln!("  {}: {} cases", t, cnt);
    }

    // Show 10 sample cases
    eprintln!("\n=== Sample cases (first 10) ===");
    for (i, c) in cases.iter().take(10).enumerate() {
        eprintln!("\n--- Case {} ---", i + 1);
        eprintln!("  Name: {}", c.name);
        eprintln!("  SELECT */alias.*: {}", c.has_select_star);
        eprintln!("  Explicit STORE_ID ref: {}", c.has_explicit_store_id);
        eprintln!(
            "  STORE_ID in schema tables: {:?}",
            c.store_id_in_schema_tables
        );
        eprintln!("  STORE_ID sources: {:?}", c.store_id_sources);
        eprintln!("  Extra cols: {:?}", c.extra_cols);
        eprintln!(
            "  Would match without extras: {}",
            c.would_match_without_extras
        );
        eprintln!("  SQL: {}", c.sql_snippet);
    }

    // Export for streamlit
    let json = serde_json::to_string_pretty(&cases).unwrap();
    std::fs::write("/tmp/store_id_cases.json", &json).unwrap();
    eprintln!(
        "\nExported {} cases to /tmp/store_id_cases.json",
        cases.len()
    );
}
