//! Test different filtering strategies on superset queries to find the highest-impact fix.
//!
//! Strategies:
//! A) Remove _HT_OP_TYPE from source values
//! B) Restrict output columns to expected keys only
//! C) Both A + B
//! D) Remove ALL same-table extras (not just _HT_OP_TYPE)
//!
//! For each strategy, count how many queries would become 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 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()
    }
}

/// Strategy A: Remove _HT_OP_TYPE from source values
fn apply_strategy_a(actual: &HashMap<String, Vec<String>>) -> HashMap<String, Vec<String>> {
    actual
        .iter()
        .map(|(k, v)| {
            let filtered: Vec<String> = v
                .iter()
                .filter(|s| {
                    let upper = s.to_uppercase();
                    !upper.ends_with("\"_HT_OP_TYPE\"")
                })
                .cloned()
                .collect();
            (k.clone(), filtered)
        })
        .collect()
}

/// Strategy B: Restrict output columns to expected keys only
fn apply_strategy_b(
    actual: &HashMap<String, Vec<String>>,
    expected: &HashMap<String, Vec<String>>,
) -> HashMap<String, Vec<String>> {
    actual
        .iter()
        .filter(|(k, _)| expected.contains_key(k.as_str()))
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect()
}

/// Strategy D: Remove same-table extras (keep only sources whose table matches expected)
fn apply_strategy_d(
    actual: &HashMap<String, Vec<String>>,
    expected: &HashMap<String, Vec<String>>,
) -> HashMap<String, Vec<String>> {
    actual
        .iter()
        .map(|(k, actual_sources)| {
            if let Some(exp_sources) = expected.get(k) {
                let expected_tables: std::collections::HashSet<String> =
                    exp_sources.iter().map(|s| table_of(s)).collect();
                let filtered: Vec<String> = actual_sources
                    .iter()
                    .filter(|s| {
                        let table = table_of(s);
                        // Keep if: table is in expected tables, or table not found
                        table.is_empty() || expected_tables.contains(&table)
                    })
                    .cloned()
                    .collect();
                (k.clone(), filtered)
            } else {
                (k.clone(), actual_sources.clone())
            }
        })
        .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_baseline = 0u64;
    let mut exact_a = 0u64; // filter _HT_OP_TYPE sources
    let mut exact_b = 0u64; // restrict to expected keys
    let mut exact_c = 0u64; // A + B
    let mut exact_d = 0u64; // remove same-table extras
    let mut exact_db = 0u64; // D + B (same-table extras + restrict keys)

    // Track what's causing the gap between B and exact
    let mut b_gained_but_not_baseline = 0u64;
    let mut c_gained_but_not_b = 0u64;

    // Track polyglot-only (no supplement/expand) to see where extras come from
    let mut extra_keys_only = 0u64; // queries where restricting keys alone fixes it
    let mut extra_sources_only = 0u64; // queries where filtering sources alone fixes it

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

        let is_exact = actual == expected;
        if is_exact {
            exact_baseline += 1;
            exact_a += 1;
            exact_b += 1;
            exact_c += 1;
            exact_d += 1;
            exact_db += 1;
            continue;
        }

        // Apply strategies
        let after_a = normalize(&apply_strategy_a(&actual));
        let after_b = normalize(&apply_strategy_b(&actual, &expected));
        let after_c = normalize(&apply_strategy_b(&after_a, &expected));
        let after_d = normalize(&apply_strategy_d(&actual, &expected));
        let after_db = normalize(&apply_strategy_b(&after_d, &expected));

        if after_a == expected {
            exact_a += 1;
        }
        if after_b == expected {
            exact_b += 1;
            b_gained_but_not_baseline += 1;
        }
        if after_c == expected {
            exact_c += 1;
            if after_b != expected {
                c_gained_but_not_b += 1;
            }
        }
        if after_d == expected {
            exact_d += 1;
        }
        if after_db == expected {
            exact_db += 1;
        }

        // Categorize
        if after_b == expected && after_a != expected {
            extra_keys_only += 1;
        }
        if after_a == expected && after_b != expected {
            extra_sources_only += 1;
        }
    }

    println!("Total:                 {total}");
    println!();
    println!("=== Exact Match Counts ===");
    println!(
        "Baseline (no filter):  {exact_baseline} ({:.1}%)",
        exact_baseline as f64 / total as f64 * 100.0
    );
    println!(
        "A) Filter _HT_OP_TYPE: {exact_a} ({:.1}%)",
        exact_a as f64 / total as f64 * 100.0
    );
    println!(
        "B) Restrict keys:      {exact_b} ({:.1}%)",
        exact_b as f64 / total as f64 * 100.0
    );
    println!(
        "C) A + B:              {exact_c} ({:.1}%)",
        exact_c as f64 / total as f64 * 100.0
    );
    println!(
        "D) Remove same-tbl:    {exact_d} ({:.1}%)",
        exact_d as f64 / total as f64 * 100.0
    );
    println!(
        "D+B) Same-tbl + keys:  {exact_db} ({:.1}%)",
        exact_db as f64 / total as f64 * 100.0
    );
    println!();
    println!("=== Gains ===");
    println!(
        "A gains over baseline: {} queries",
        exact_a - exact_baseline
    );
    println!(
        "B gains over baseline: {} queries (extra output column keys only)",
        b_gained_but_not_baseline
    );
    println!(
        "C gains over B:        {} queries (source filtering on top of key restriction)",
        c_gained_but_not_b
    );
    println!(
        "D+B gains over B:      {} queries (same-table source filtering)",
        exact_db - exact_b
    );
    println!();
    println!("=== Issue Breakdown ===");
    println!(
        "Extra keys only:       {} (restricting keys fixes it, source filtering doesn't)",
        extra_keys_only
    );
    println!(
        "Extra sources only:    {} (source filtering fixes it, key restriction doesn't)",
        extra_sources_only
    );
}
