//! Test UNION ALL handling in the column lineage engine.
//!
//! Runs several UNION ALL query patterns through `column_lineage` and prints
//! `column_dict` for each, sorted for deterministic output.
//!
//! Run: cargo run --example test_union --release

use altimate_core::lineage::{ColumnLineageOptions, column_lineage};
use altimate_core::schema::types::SqlDialect;
use std::collections::HashMap;

/// Build a minimal schema map: table_name -> JSON array of column names.
fn schema(tables: &[(&str, &[&str])]) -> HashMap<String, serde_json::Value> {
    tables
        .iter()
        .map(|(name, cols)| {
            let columns: Vec<serde_json::Value> = cols
                .iter()
                .map(|c| serde_json::json!({"name": c, "type": "VARCHAR"}))
                .collect();
            (name.to_string(), serde_json::json!(columns))
        })
        .collect()
}

fn run_test(label: &str, sql: &str, tables: &[(&str, &[&str])]) {
    println!("{}", "=".repeat(70));
    println!("Test: {label}");
    println!("SQL:  {sql}");
    println!("{}", "-".repeat(70));

    let s = schema(tables);
    let opts = ColumnLineageOptions {
        inject_shared: true,
        ..Default::default()
    };

    match column_lineage(sql, SqlDialect::Snowflake, &s, Some(&opts)) {
        Ok(result) => {
            let mut entries: Vec<_> = result.column_dict.iter().collect();
            entries.sort_by(|a, b| a.0.cmp(b.0));

            println!("column_dict ({} columns):", entries.len());
            for (col, sources) in &entries {
                let mut sorted_sources = (*sources).clone();
                sorted_sources.sort();
                println!("  {col}: {sorted_sources:?}");
            }

            if !result.errors.is_empty() {
                println!("errors: {:?}", result.errors);
            }
        }
        Err(e) => {
            println!("ERROR: {e}");
        }
    }

    println!();
}

fn main() {
    // 1. Simple UNION ALL
    run_test(
        "1. Simple UNION ALL",
        "SELECT a, b FROM table1 UNION ALL SELECT c, d FROM table2",
        &[
            ("table1", &["a", "b", "c"][..]),
            ("table2", &["c", "d", "e"][..]),
        ],
    );

    // 2. UNION ALL with aggregate
    run_test(
        "2. UNION ALL with aggregate",
        "SELECT a, MAX(b) AS max_b FROM table1 GROUP BY a \
         UNION ALL \
         SELECT c, SUM(d) AS max_b FROM table2 GROUP BY c",
        &[
            ("table1", &["a", "b", "c"][..]),
            ("table2", &["c", "d", "e"][..]),
        ],
    );

    // 3. UNION ALL in subquery
    run_test(
        "3. UNION ALL in subquery",
        "SELECT x, y FROM (\
            SELECT a AS x, b AS y FROM table1 \
            UNION ALL \
            SELECT c AS x, d AS y FROM table2\
         ) t",
        &[
            ("table1", &["a", "b", "c"][..]),
            ("table2", &["c", "d", "e"][..]),
        ],
    );

    // 4. UNION ALL with star
    run_test(
        "4. UNION ALL with star",
        "SELECT * FROM table1 UNION ALL SELECT * FROM table2",
        &[
            ("table1", &["a", "b", "c"][..]),
            ("table2", &["a", "b", "c"][..]),
        ],
    );

    // 5. Star in FROM subquery UNION
    run_test(
        "5. Star in FROM subquery UNION",
        "SELECT * FROM (\
            SELECT * FROM table1 \
            UNION ALL \
            SELECT * FROM table2\
         ) t",
        &[
            ("table1", &["a", "b", "c"][..]),
            ("table2", &["a", "b", "c"][..]),
        ],
    );

    // 6. Mixed star + named in UNION branch
    run_test(
        "6. Mixed star + named in UNION branch",
        "SELECT a, b, c FROM table1 UNION ALL SELECT * FROM table2",
        &[
            ("table1", &["a", "b", "c"][..]),
            ("table2", &["a", "b", "c"][..]),
        ],
    );

    // 7. Three-branch UNION
    run_test(
        "7. Three-branch UNION",
        "SELECT a, b FROM table1 \
         UNION ALL SELECT a, b FROM table2 \
         UNION ALL SELECT a, b FROM table3",
        &[
            ("table1", &["a", "b"][..]),
            ("table2", &["a", "b"][..]),
            ("table3", &["a", "b"][..]),
        ],
    );
}
