use altimate_core::lineage::column_lineage;
use altimate_core::lineage::polyglot_column_lineage;
use altimate_core::schema::types::SqlDialect;
use std::collections::HashMap;

fn run(label: &str, sql: &str, schema: &HashMap<String, serde_json::Value>, dialect: SqlDialect) {
    println!("--- {label} ---");
    match column_lineage(sql, dialect, schema, None) {
        Ok(r) => {
            if r.column_dict.is_empty() {
                println!("  (empty column_dict)");
            }
            let mut keys: Vec<_> = r.column_dict.keys().collect();
            keys.sort();
            for k in keys {
                println!("  {k}: {:?}", r.column_dict[k]);
            }
            if !r.errors.is_empty() {
                println!("  ERRORS: {:?}", r.errors);
            }
        }
        Err(e) => println!("  ERROR: {e}"),
    }
    println!();
}

fn main() {
    let schema: HashMap<String, serde_json::Value> = serde_json::from_str(r#"{
        "table1": {"a": "int", "b": "int", "c": "int", "d": "int", "e": "date"},
        "table2": {"a": "int", "b": "int", "c": "int", "d": "int", "e": "date", "x": "int", "y": "int"}
    }"#).unwrap();

    // 1. with_stars — qualified wildcard t1.*
    let sql = r#"WITH cte1 AS (SELECT * FROM table1 WHERE func(d) > e AND e > CURRENT_DATE),
cte2 AS (SELECT * FROM table2)
SELECT t1.*, t2.x, t2.y FROM cte1 t1 JOIN cte2 t2 ON t1.a = t2.a AND t1.b = t2.b"#;
    run("with_stars bigquery", sql, &schema, SqlDialect::BigQuery);
    run("with_stars postgres", sql, &schema, SqlDialect::Postgres);

    // 2. with_rank — CTE expression column
    let sql2 = "WITH cte AS (SELECT a, b, RANK() OVER (PARTITION BY a, b ORDER BY c) AS rank FROM table1) SELECT a, b FROM cte WHERE rank = 1";
    run("with_rank bigquery", sql2, &schema, SqlDialect::BigQuery);
    run("with_rank postgres", sql2, &schema, SqlDialect::Postgres);

    // 3. with_subquery — IN subquery
    let sql3 = "SELECT a, b FROM table1 WHERE e > CURRENT_DATE AND a IN (SELECT t2.a FROM table2 AS t2 WHERE t2.b > 10)";
    run(
        "with_subquery bigquery",
        sql3,
        &schema,
        SqlDialect::BigQuery,
    );
    run(
        "with_subquery postgres",
        sql3,
        &schema,
        SqlDialect::Postgres,
    );

    // 4. subquery_in_for_clause — FROM subquery with alias
    let sql4 = "SELECT t1.a, t1.b FROM (SELECT table1.* FROM table1 WHERE table1.e > CURRENT_DATE) AS t1 WHERE t1.a > 10";
    run(
        "subquery_in_for bigquery",
        sql4,
        &schema,
        SqlDialect::BigQuery,
    );
    run(
        "subquery_in_for postgres",
        sql4,
        &schema,
        SqlDialect::Postgres,
    );

    // 5. lateral_joins — LATERAL correlation
    let sql5 = "SELECT a, b, c FROM table1, LATERAL (SELECT a, b, c FROM table2 WHERE table1.a = table2.a) AS table2";
    run(
        "lateral_joins bigquery",
        sql5,
        &schema,
        SqlDialect::BigQuery,
    );

    // 6. lateral_view_flatten
    let schema_nested: HashMap<String, serde_json::Value> = serde_json::from_str(
        r#"{
        "db": {"schema": {"table1": {"a": "INTEGER", "b": "INTEGER", "c": "VARIANT"}}}
    }"#,
    )
    .unwrap();
    let sql6a = "SELECT a, b, om.value as c FROM (SELECT a, b, c FROM db.schema.table1)t, lateral flatten(input => t.c) om";
    // Check parser behavior
    use altimate_core::validator::dialect::get_parser_dialect;
    use datafusion::sql::sqlparser::ast::SetExpr;
    use datafusion::sql::sqlparser::parser::Parser;
    for d in &[
        SqlDialect::Postgres,
        SqlDialect::BigQuery,
        SqlDialect::Snowflake,
    ] {
        let pd = get_parser_dialect(*d);
        match Parser::parse_sql(pd.as_ref(), sql6a) {
            Ok(stmts) => println!(
                "  {:?} parse OK: {} stmts, from_len={}",
                d,
                stmts.len(),
                if let datafusion::sql::sqlparser::ast::Statement::Query(q) = &stmts[0] {
                    if let SetExpr::Select(sel) = &*q.body {
                        format!("{} items, {} from", sel.projection.len(), sel.from.len())
                    } else {
                        "non-select".to_string()
                    }
                } else {
                    "non-query".to_string()
                }
            ),
            Err(e) => println!("  {:?} parse FAIL: {}", d, e),
        }
    }
    println!();

    println!("--- lateral_view_flatten RAW EDGES ---");
    match polyglot_column_lineage(sql6a, SqlDialect::BigQuery) {
        Ok(r) => {
            for e in &r.edges {
                println!(
                    "  src={}:{} -> tgt={} (kind={:?})",
                    e.source_table, e.source_column, e.target_column, e.terminal_kind
                );
            }
            if !r.errors.is_empty() {
                println!("  ERRORS: {:?}", r.errors);
            }
        }
        Err(e) => println!("  ERROR: {e}"),
    }
    println!();
    run(
        "lateral_view_flatten bigquery",
        sql6a,
        &schema_nested,
        SqlDialect::BigQuery,
    );

    // 7. multiple_cte_and_joins_unions_with_stars — UNION ALL with star
    let sql6 = "WITH cte1 AS (SELECT a, b, c FROM table1 WHERE func(d) > e), cte2 AS (SELECT a, b, c FROM table2) SELECT * FROM cte1 UNION ALL SELECT * FROM cte2";
    println!("--- multi_cte_union_stars snowflake RAW EDGES ---");
    match polyglot_column_lineage(sql6, SqlDialect::Snowflake) {
        Ok(r) => {
            for e in &r.edges {
                println!(
                    "  src={}:{} -> tgt={} (kind={:?})",
                    e.source_table, e.source_column, e.target_column, e.terminal_kind
                );
            }
            if !r.errors.is_empty() {
                println!("  ERRORS: {:?}", r.errors);
            }
        }
        Err(e) => println!("  ERROR: {e}"),
    }
    println!();
    run(
        "multi_cte_union_stars snowflake",
        sql6,
        &schema,
        SqlDialect::Snowflake,
    );
    run(
        "multi_cte_union_stars bigquery",
        sql6,
        &schema,
        SqlDialect::BigQuery,
    );
}
