//! Categorize WHY extra output columns appear in column lineage results.
//!
//! For each query where we produce columns not in the expected set,
//! classifies each extra column into one of:
//!   - `window_function`: name matches ROW_NUMBER, RANK, etc.
//!   - `select_star`: SQL contains SELECT * or alias.*
//!   - `literal_alias`: column has empty sources (came from a literal)
//!   - `expression_alias`: column has sources (computed expression)
//!   - `hightouch_metadata`: column name starts with _HT_
//!   - `positional_col`: column name like _col_N / _COL_N
//!   - `other`: everything else
//!
//! Run: cargo run --example extra_cols_analysis --release

use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use serde::Deserialize;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::BufRead;
use std::path::Path;

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

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

/// Category explaining why an extra column appeared.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
enum ExtraCategory {
    WindowFunction,
    SelectStar,
    LiteralAlias,
    ExpressionAlias,
    HightouchMetadata,
    PositionalCol,
    Other,
}

impl ExtraCategory {
    fn label(self) -> &'static str {
        match self {
            Self::WindowFunction => "window_function",
            Self::SelectStar => "select_star",
            Self::LiteralAlias => "literal_alias",
            Self::ExpressionAlias => "expression_alias",
            Self::HightouchMetadata => "hightouch_metadata",
            Self::PositionalCol => "positional_col",
            Self::Other => "other",
        }
    }
}

/// One sample of an extra column for display purposes.
struct CategorySample {
    query_name: String,
    sql_snippet: String,
    extra_col: String,
    sources: Vec<String>,
}

// ---------------------------------------------------------------------------
// Known window function output names (uppercased for matching)
// ---------------------------------------------------------------------------

const WINDOW_FUNCTION_NAMES: &[&str] = &[
    "ROW_NUMBER",
    "RANK",
    "DENSE_RANK",
    "NTILE",
    "CUME_DIST",
    "PERCENT_RANK",
    "LAG",
    "LEAD",
    "FIRST_VALUE",
    "LAST_VALUE",
    "NTH_VALUE",
    "ROW_NUM",
    "ROWNUM",
    "RN",
    "RNUM",
    "ROWNO",
    "ROW_NO",
];

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

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 sql_snippet(sql: &str, max_len: usize) -> String {
    let flat: String = sql
        .chars()
        .map(|c| if c == '\n' || c == '\r' { ' ' } else { c })
        .collect();
    if flat.len() > max_len {
        format!("{}...", &flat[..max_len])
    } else {
        flat
    }
}

fn has_select_star(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    upper.contains("SELECT *") || upper.contains("SELECT\n*") || upper.contains("SELECT\r\n*")
}

fn has_alias_dot_star(sql: &str) -> bool {
    let upper = sql.to_uppercase();
    for (i, _) in upper.match_indices(".*") {
        if i > 0 {
            let prev = sql.as_bytes().get(i - 1).copied().unwrap_or(b' ');
            if prev.is_ascii_alphanumeric() || prev == b'_' || prev == b'"' {
                return true;
            }
        }
    }
    false
}

/// Classify a single extra column into a category.
fn classify_extra_column(col_name: &str, sources: &[String], sql: &str) -> ExtraCategory {
    let upper_name = col_name.trim_matches('"').to_uppercase();

    // 1. Hightouch metadata columns
    if upper_name.starts_with("_HT_") {
        return ExtraCategory::HightouchMetadata;
    }

    // 2. Positional columns like _col_0, _COL_1, _C0, etc.
    if upper_name.starts_with("_COL_") || upper_name.starts_with("_C") {
        let suffix = if upper_name.starts_with("_COL_") {
            &upper_name[5..]
        } else {
            &upper_name[2..]
        };
        if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) {
            return ExtraCategory::PositionalCol;
        }
    }

    // 3. Window function output names
    for wf in WINDOW_FUNCTION_NAMES {
        if upper_name == *wf {
            return ExtraCategory::WindowFunction;
        }
    }

    // 4. SELECT * or alias.* expansion
    if has_select_star(sql) || has_alias_dot_star(sql) {
        return ExtraCategory::SelectStar;
    }

    // 5. Literal alias (empty sources = came from a literal/constant)
    if sources.is_empty() {
        return ExtraCategory::LiteralAlias;
    }

    // 6. Expression alias (has sources = computed expression not in expected)
    if !sources.is_empty() {
        return ExtraCategory::ExpressionAlias;
    }

    ExtraCategory::Other
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

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

    let mut total = 0u64;
    let mut parse_fail = 0u64;
    let mut exact_match = 0u64;
    let mut queries_with_extras = 0u64;
    let mut total_extra_cols = 0u64;

    // Extra column name frequency (for top-20 listing).
    let mut extra_col_freq: BTreeMap<String, u64> = BTreeMap::new();

    // Per-category counters and samples.
    let mut category_counts: BTreeMap<ExtraCategory, u64> = BTreeMap::new();
    let mut category_query_counts: BTreeMap<ExtraCategory, HashSet<u64>> = BTreeMap::new();
    let mut category_samples: BTreeMap<ExtraCategory, Vec<CategorySample>> = BTreeMap::new();

    let mut query_idx = 0u64;

    for line in reader.lines() {
        let line = line.expect("failed to read line");
        let line = line.trim().to_string();
        if line.is_empty() {
            continue;
        }

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

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

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

        if actual == expected {
            exact_match += 1;
            continue;
        }

        let actual_keys: HashSet<&String> = actual.keys().collect();
        let expected_keys: HashSet<&String> = expected.keys().collect();

        let extra: Vec<String> = actual_keys
            .difference(&expected_keys)
            .map(|s| (*s).clone())
            .collect();

        if extra.is_empty() {
            continue;
        }

        queries_with_extras += 1;
        total_extra_cols += extra.len() as u64;

        for col in &extra {
            *extra_col_freq.entry(col.clone()).or_default() += 1;

            let sources = actual.get(col).cloned().unwrap_or_default();
            let category = classify_extra_column(col, &sources, &case.sql);

            *category_counts.entry(category).or_default() += 1;
            category_query_counts
                .entry(category)
                .or_default()
                .insert(query_idx);

            let samples = category_samples.entry(category).or_default();
            if samples.len() < 3 {
                samples.push(CategorySample {
                    query_name: case.name.clone(),
                    sql_snippet: sql_snippet(&case.sql, 250),
                    extra_col: col.clone(),
                    sources,
                });
            }
        }
    }

    // =====================================================================
    // Output
    // =====================================================================
    let sep = "=".repeat(80);
    let dash = "-".repeat(60);

    // ── Summary ──────────────────────────────────────────────────────────
    println!("\n{sep}");
    println!("  EXTRA COLUMNS ANALYSIS — WHY DO THEY APPEAR?");
    println!("{sep}");
    println!("Total queries processed:      {total}");
    println!("Parse failures:               {parse_fail}");
    println!("Exact matches:                {exact_match}");
    println!("Queries with extra columns:   {queries_with_extras}");
    println!("Total extra column instances:  {total_extra_cols}");

    // ── Top 20 extra column names ────────────────────────────────────────
    println!("\n{dash}");
    println!("  Top 20 Extra Column Names");
    println!("{dash}");
    let mut freq_vec: Vec<_> = extra_col_freq.iter().collect();
    freq_vec.sort_by(|a, b| b.1.cmp(a.1));
    for (i, (col, count)) in freq_vec.iter().take(20).enumerate() {
        println!("  {:>3}. {:<50} {:>6}", i + 1, col, count);
    }

    // ── Counts per category ──────────────────────────────────────────────
    println!("\n{dash}");
    println!("  Extra Columns by Category");
    println!("{dash}");
    println!("  {:<25} {:>10} {:>12}", "Category", "Columns", "Queries");
    println!("  {}", "-".repeat(50));

    let all_categories = [
        ExtraCategory::WindowFunction,
        ExtraCategory::SelectStar,
        ExtraCategory::LiteralAlias,
        ExtraCategory::ExpressionAlias,
        ExtraCategory::HightouchMetadata,
        ExtraCategory::PositionalCol,
        ExtraCategory::Other,
    ];

    for cat in &all_categories {
        let col_count = category_counts.get(cat).copied().unwrap_or(0);
        let q_count = category_query_counts.get(cat).map(|s| s.len()).unwrap_or(0);
        println!("  {:<25} {:>10} {:>12}", cat.label(), col_count, q_count);
    }

    // ── Samples per category ─────────────────────────────────────────────
    for cat in &all_categories {
        let samples = match category_samples.get(cat) {
            Some(s) if !s.is_empty() => s,
            _ => continue,
        };

        println!("\n{dash}");
        println!(
            "  Samples: {} ({} total columns)",
            cat.label(),
            category_counts.get(cat).copied().unwrap_or(0)
        );
        println!("{dash}");

        for (i, sample) in samples.iter().take(3).enumerate() {
            println!("\n  Sample {}:", i + 1);
            println!("    Query:       {}", sample.query_name);
            println!("    Extra col:   {}", sample.extra_col);
            println!("    Sources:     {:?}", sample.sources);
            println!("    SQL:         {}", sample.sql_snippet);
        }
    }

    println!("\n{sep}");
    println!("  DONE");
    println!("{sep}\n");
}
