//! Analyze extra output column keys (columns we produce that ground truth doesn't expect).
//! These are the columns causing 707 queries to be superset instead of exact.

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(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;

    // Count frequency of extra column names
    let mut extra_col_freq: HashMap<String, u64> = HashMap::new();
    let mut total_queries_with_extras = 0u64;
    let mut total_extra_keys = 0u64;

    // Track queries where restricting keys would make them exact
    let mut fixable_extra_col_freq: HashMap<String, u64> = HashMap::new();
    let mut fixable_count = 0u64;

    // Sample: queries where extra keys are the only problem
    let mut samples: Vec<(String, Vec<String>, Vec<String>)> = Vec::new(); // (sql_prefix, extra_keys, expected_keys)

    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,
        };

        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(&result.column_dict);
        let expected = normalize(&case.expected);

        if actual == expected {
            continue;
        }

        // Find extra keys
        let extra_keys: Vec<String> = actual
            .keys()
            .filter(|k| !expected.contains_key(k.as_str()))
            .cloned()
            .collect();

        if extra_keys.is_empty() {
            continue;
        }

        total_queries_with_extras += 1;
        total_extra_keys += extra_keys.len() as u64;

        for k in &extra_keys {
            *extra_col_freq.entry(k.clone()).or_default() += 1;
        }

        // Check if restricting to expected keys makes it exact
        let restricted: HashMap<String, Vec<String>> = actual
            .iter()
            .filter(|(k, _)| expected.contains_key(k.as_str()))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        if restricted == expected {
            fixable_count += 1;
            for k in &extra_keys {
                *fixable_extra_col_freq.entry(k.clone()).or_default() += 1;
            }
            if samples.len() < 5 {
                let sql_prefix: String = case.sql.chars().take(120).collect();
                let expected_keys: Vec<String> = expected.keys().cloned().collect();
                samples.push((sql_prefix, extra_keys.clone(), expected_keys));
            }
        }
    }

    println!("=== Extra Output Column Key Analysis ===");
    println!("Queries with extra keys:     {total_queries_with_extras}");
    println!("Total extra keys:            {total_extra_keys}");
    println!("Fixable (restrict→exact):    {fixable_count}");
    println!();

    println!("=== Top 20 Extra Column Names (all queries) ===");
    let mut freq_vec: Vec<_> = extra_col_freq.into_iter().collect();
    freq_vec.sort_by(|a, b| b.1.cmp(&a.1));
    for (col, count) in freq_vec.iter().take(20) {
        println!("  {col:<40} {count:>6}");
    }

    println!();
    println!("=== Top 20 Extra Column Names (fixable queries only) ===");
    let mut fix_vec: Vec<_> = fixable_extra_col_freq.into_iter().collect();
    fix_vec.sort_by(|a, b| b.1.cmp(&a.1));
    for (col, count) in fix_vec.iter().take(20) {
        println!("  {col:<40} {count:>6}");
    }

    println!();
    println!("=== Sample Fixable Queries ===");
    for (i, (sql, extras, expected_keys)) in samples.iter().enumerate() {
        println!("--- Sample {} ---", i + 1);
        println!("SQL:           {}...", sql);
        println!("Extra keys:    {:?}", extras);
        println!("Expected keys: {:?}", expected_keys);
        println!();
    }
}
