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;

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

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()
}

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("open");
    let reader = std::io::BufReader::with_capacity(1024 * 1024, file);
    let dialect = SqlDialect::Snowflake;

    let mut total = 0u64;
    let mut exact_before = 0u64;
    let mut exact_after = 0u64;

    for line in reader.lines() {
        let line = line.unwrap();
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let case: Case = match serde_json::from_str(line) {
            Ok(c) => c,
            Err(_) => continue,
        };
        total += 1;

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

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

        if actual == expected {
            exact_before += 1;
            exact_after += 1;
            continue;
        }

        // Filter out _HT_OP_TYPE source references
        let filtered: HashMap<String, Vec<String>> = actual
            .iter()
            .map(|(k, v)| {
                let filtered_v: Vec<String> = v
                    .iter()
                    .filter(|s| {
                        let upper = s.to_uppercase();
                        !upper.ends_with("\"_HT_OP_TYPE\"")
                    })
                    .cloned()
                    .collect();
                (k.clone(), filtered_v)
            })
            .collect();

        if filtered == expected {
            exact_after += 1;
        }
    }

    println!("Total:        {total}");
    println!(
        "Exact before: {exact_before} ({:.1}%)",
        exact_before as f64 / total as f64 * 100.0
    );
    println!(
        "Exact after:  {exact_after} ({:.1}%)",
        exact_after as f64 / total as f64 * 100.0
    );
    println!("Gained:       {}", exact_after - exact_before);
}
