//! Export all LATERAL FLATTEN cases from the ground truth benchmark to JSON
//! for Streamlit review.
//!
//! Run: `cargo run --release --example flatten_cases_export`
//! Output: `/tmp/flatten_cases.json`

use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use std::collections::{BTreeMap, 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 FlattenCase {
    name: String,
    sql: String,
    schema_tables: Vec<String>,
    columns: Vec<ColumnComparison>,
    status: String, // "exact_match", "extra_only", "missing_only", "mixed"
}

#[derive(serde::Serialize)]
struct ColumnComparison {
    column: String,
    expected: Vec<String>,
    actual: Vec<String>,
    extra: Vec<String>,
    missing: Vec<String>,
    status: String, // "match", "extra", "missing", "mixed"
}

fn schema_tables(schema: &serde_json::Value) -> Vec<String> {
    let mut tables = Vec::new();
    if let Some(obj) = schema.as_object() {
        for (k, v) in obj {
            if v.is_object() && v.as_object().unwrap().values().all(|vv| vv.is_string()) {
                tables.push(k.to_uppercase());
            } else if let Some(inner) = v.as_object() {
                for (k2, v2) in inner {
                    if v2.is_object() && v2.as_object().unwrap().values().all(|vv| vv.is_string()) {
                        tables.push(format!("{}.{}", k.to_uppercase(), k2.to_uppercase()));
                    } else if let Some(inner2) = v2.as_object() {
                        for (k3, v3) in inner2 {
                            if v3.is_object() {
                                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 = std::path::Path::new("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 flatten_cases: Vec<FlattenCase> = Vec::new();
    let mut total = 0;
    let mut flatten_count = 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,
        };

        // Check if SQL contains LATERAL FLATTEN
        let sql_upper = case.sql.to_uppercase();
        if !sql_upper.contains("FLATTEN") {
            continue;
        }
        flatten_count += 1;

        // Run our engine
        let schema_map = to_schema_map(&case.schema);
        let result = column_lineage(&case.sql, SqlDialect::Snowflake, &schema_map, None);

        let actual_dict: BTreeMap<String, BTreeSet<String>> = match result {
            Ok(r) => r
                .column_dict
                .into_iter()
                .map(|(k, v)| (k, v.into_iter().collect::<BTreeSet<String>>()))
                .collect(),
            Err(_) => BTreeMap::new(),
        };

        // Compare columns
        let mut all_cols: BTreeSet<String> = BTreeSet::new();
        for k in case.expected.keys() {
            all_cols.insert(k.to_uppercase());
        }
        for k in actual_dict.keys() {
            all_cols.insert(k.to_uppercase());
        }

        let mut columns = Vec::new();
        let mut has_extra_col = false;
        let mut has_missing_col = false;
        let mut has_source_diff = false;

        for col in &all_cols {
            let expected_sources: BTreeSet<String> = case
                .expected
                .iter()
                .find(|(k, _)| k.to_uppercase() == *col)
                .map(|(_, v)| v.iter().map(|s| s.to_uppercase()).collect())
                .unwrap_or_default();

            let actual_sources: BTreeSet<String> = actual_dict
                .iter()
                .find(|(k, _)| k.to_uppercase() == *col)
                .map(|(_, v)| v.iter().map(|s| s.to_uppercase()).collect())
                .unwrap_or_default();

            let in_expected = case.expected.keys().any(|k| k.to_uppercase() == *col);
            let in_actual = actual_dict.keys().any(|k| k.to_uppercase() == *col);

            let extra: Vec<String> = actual_sources
                .difference(&expected_sources)
                .cloned()
                .collect();
            let missing: Vec<String> = expected_sources
                .difference(&actual_sources)
                .cloned()
                .collect();

            let status = if !in_expected {
                has_extra_col = true;
                "extra_col".to_string()
            } else if !in_actual {
                has_missing_col = true;
                "missing_col".to_string()
            } else if extra.is_empty() && missing.is_empty() {
                "match".to_string()
            } else if !extra.is_empty() && missing.is_empty() {
                has_source_diff = true;
                "extra_sources".to_string()
            } else if extra.is_empty() && !missing.is_empty() {
                has_source_diff = true;
                "missing_sources".to_string()
            } else {
                has_source_diff = true;
                "mixed".to_string()
            };

            columns.push(ColumnComparison {
                column: col.clone(),
                expected: expected_sources.into_iter().collect(),
                actual: actual_sources.into_iter().collect(),
                extra,
                missing,
                status,
            });
        }

        let overall = if !has_extra_col && !has_missing_col && !has_source_diff {
            "exact_match"
        } else if has_extra_col && !has_missing_col && !has_source_diff {
            "extra_cols_only"
        } else if !has_extra_col && has_missing_col && !has_source_diff {
            "missing_cols_only"
        } else {
            "mixed"
        };

        flatten_cases.push(FlattenCase {
            name: case.name,
            sql: case.sql,
            schema_tables: schema_tables(&case.schema),
            columns,
            status: overall.to_string(),
        });
    }

    // Sort: non-matches first, then by name
    flatten_cases.sort_by(|a, b| {
        let a_match = a.status == "exact_match";
        let b_match = b.status == "exact_match";
        a_match.cmp(&b_match).then(a.name.cmp(&b.name))
    });

    let json = serde_json::to_string_pretty(&flatten_cases).unwrap();
    std::fs::write("/tmp/flatten_cases.json", &json).unwrap();

    eprintln!("Total queries: {}", total);
    eprintln!("FLATTEN queries: {}", flatten_count);
    eprintln!("Exported to /tmp/flatten_cases.json");

    // Summary
    let exact = flatten_cases
        .iter()
        .filter(|c| c.status == "exact_match")
        .count();
    let extra = flatten_cases
        .iter()
        .filter(|c| c.status == "extra_cols_only")
        .count();
    let missing = flatten_cases
        .iter()
        .filter(|c| c.status == "missing_cols_only")
        .count();
    let mixed = flatten_cases.iter().filter(|c| c.status == "mixed").count();
    eprintln!("  exact_match: {}", exact);
    eprintln!("  extra_cols_only: {}", extra);
    eprintln!("  missing_cols_only: {}", missing);
    eprintln!("  mixed: {}", mixed);
}
