use altimate_core::lineage::polyglot_column_lineage;
use altimate_core::schema::types::SqlDialect;

fn dump(label: &str, sql: &str, dialects: &[(&str, SqlDialect)]) {
    println!("=== {label} ===");
    for (dname, dialect) in dialects {
        print!("  [{dname}] ");
        match polyglot_column_lineage(sql, *dialect) {
            Ok(r) => {
                if r.edges.is_empty() && r.errors.is_empty() {
                    println!("(no edges)");
                } else {
                    println!();
                    for e in &r.edges {
                        println!(
                            "    src={}:{} -> tgt={} (kind={:?})",
                            e.source_table, e.source_column, e.target_column, e.terminal_kind
                        );
                    }
                    for e in &r.errors {
                        println!("    ERR: {}", e.message);
                    }
                }
            }
            Err(e) => println!("ERROR: {e}"),
        }
    }
    println!();
}

fn test_parse(label: &str, sql: &str) {
    use datafusion::sql::sqlparser::dialect::{
        BigQueryDialect, GenericDialect, PostgreSqlDialect, RedshiftSqlDialect, SnowflakeDialect,
    };
    use datafusion::sql::sqlparser::parser::Parser;
    let dialects: Vec<(&str, Box<dyn datafusion::sql::sqlparser::dialect::Dialect>)> = vec![
        ("generic", Box::new(GenericDialect {})),
        ("bigquery", Box::new(BigQueryDialect)),
        ("snowflake", Box::new(SnowflakeDialect)),
        ("postgres", Box::new(PostgreSqlDialect {})),
        ("redshift", Box::new(RedshiftSqlDialect {})),
    ];
    println!("=== PARSE: {label} ===");
    for (name, dialect) in &dialects {
        match Parser::parse_sql(dialect.as_ref(), sql) {
            Ok(_) => println!("  [{name}] OK"),
            Err(e) => println!("  [{name}] FAIL: {e}"),
        }
    }
    println!();
}

fn main() {
    let all = &[
        ("bq", SqlDialect::BigQuery),
        ("sf", SqlDialect::Snowflake),
        ("pg", SqlDialect::Postgres),
        ("rs", SqlDialect::Redshift),
    ];
    let pg_rs = &[("pg", SqlDialect::Postgres), ("rs", SqlDialect::Redshift)];

    dump(
        "multiple_cte_union_stars",
        r#"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"#,
        &[("sf", SqlDialect::Snowflake)],
    );

    dump(
        "with_rank",
        r#"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"#,
        all,
    );

    dump(
        "with_subquery",
        r#"SELECT a, b FROM table1 WHERE e > CURRENT_DATE AND a IN (SELECT t2.a FROM table2 AS t2 WHERE t2.b > 10)"#,
        pg_rs,
    );

    dump(
        "subquery_in_for",
        r#"SELECT t1.a, t1.b FROM (SELECT table1.* FROM table1 WHERE table1.e > CURRENT_DATE) AS t1 WHERE t1.a > 10"#,
        all,
    );

    dump(
        "lateral_joins",
        r#"SELECT a, b, c FROM table1, LATERAL (SELECT a, b, c FROM table2 WHERE table1.a = table2.a) AS table2"#,
        &[("bq", SqlDialect::BigQuery)],
    );

    dump(
        "lateral_flatten",
        r#"SELECT a, b, om.value as c FROM (SELECT a, b, c FROM db.schema.table1)t, lateral flatten(input => t.c) om"#,
        &[("bq", SqlDialect::BigQuery)],
    );

    // Test with EXACT fixture SQL (has trailing commas)
    {
        use altimate_core::lineage::column_lineage;
        use std::collections::HashMap;
        let mut schema: HashMap<String, serde_json::Value> = HashMap::new();
        schema.insert(
            "table1".to_string(),
            serde_json::json!({"a": "int", "b": "int", "c": "int", "d": "int", "e": "date"}),
        );
        schema.insert("table2".to_string(), serde_json::json!({"a": "int", "b": "int", "c": "int", "d": "int", "e": "date", "x": "int", "y": "int"}));

        let sql_fixture = "WITH cte AS (\n    SELECT\n        a,\n        b,\n        RANK() OVER (PARTITION BY a, b ORDER BY c) AS rank\n    FROM\n        table1\n)\nSELECT\n    a,\n    b,\nFROM cte\nWHERE rank = 1";

        println!("=== FIXTURE with_rank [postgres] ===");
        match column_lineage(sql_fixture, SqlDialect::Postgres, &schema, None) {
            Ok(r) => {
                for (k, v) in &r.column_dict {
                    println!("  {k}: {v:?}");
                }
            }
            Err(e) => println!("  ERROR: {e}"),
        }
    }

    // Direct lineage test for postgres with_rank
    {
        use altimate_core::lineage::column_lineage;
        use std::collections::HashMap;
        let mut schema: HashMap<String, serde_json::Value> = HashMap::new();
        schema.insert(
            "table1".to_string(),
            serde_json::json!({"a": "int", "b": "int", "c": "int", "d": "int", "e": "date"}),
        );
        schema.insert("table2".to_string(), serde_json::json!({"a": "int", "b": "int", "c": "int", "d": "int", "e": "date", "x": "int", "y": "int"}));
        let sql = r#"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"#;

        for (dname, dialect) in &[
            ("postgres", SqlDialect::Postgres),
            ("bigquery", SqlDialect::BigQuery),
        ] {
            println!("=== complete_column_lineage with_rank [{dname}] ===");
            match column_lineage(sql, *dialect, &schema, None) {
                Ok(r) => {
                    for (k, v) in &r.column_dict {
                        println!("  {k}: {v:?}");
                    }
                    if !r.errors.is_empty() {
                        println!("  errors: {:?}", r.errors);
                    }
                }
                Err(e) => println!("  ERROR: {e}"),
            }
        }

        let sql2 = r#"SELECT a, b FROM table1 WHERE e > CURRENT_DATE AND a IN (SELECT t2.a FROM table2 AS t2 WHERE t2.b > 10)"#;
        for (dname, dialect) in &[
            ("postgres", SqlDialect::Postgres),
            ("bigquery", SqlDialect::BigQuery),
        ] {
            println!("=== complete_column_lineage with_subquery [{dname}] ===");
            match column_lineage(sql2, *dialect, &schema, None) {
                Ok(r) => {
                    for (k, v) in &r.column_dict {
                        println!("  {k}: {v:?}");
                    }
                }
                Err(e) => println!("  ERROR: {e}"),
            }
        }
    }

    // AST debug for lateral flatten
    {
        use datafusion::sql::sqlparser::dialect::BigQueryDialect;
        use datafusion::sql::sqlparser::parser::Parser;
        let sql = r#"SELECT a, b, om.value as c FROM (SELECT a, b, c FROM db.schema.table1)t, lateral flatten(input => t.c) om"#;
        match Parser::parse_sql(&BigQueryDialect, sql) {
            Ok(stmts) => {
                for stmt in &stmts {
                    println!("=== LATERAL FLATTEN AST ===");
                    println!("{stmt:#?}");
                }
            }
            Err(e) => println!("PARSE ERROR: {e}"),
        }
    }

    // Test parsing with different dialects — note trailing comma like fixture
    test_parse(
        "with_rank_trailing_comma",
        "WITH cte AS (\n    SELECT\n        a,\n        b,\n        RANK() OVER (PARTITION BY a, b ORDER BY c) AS rank\n    FROM\n        table1\n)\nSELECT\n    a,\n    b,\nFROM cte\nWHERE rank = 1",
    );

    test_parse(
        "with_subquery",
        r#"SELECT a, b FROM table1 WHERE e > CURRENT_DATE AND a IN (SELECT t2.a FROM table2 AS t2 WHERE t2.b > 10)"#,
    );

    test_parse(
        "subquery_in_for",
        r#"SELECT t1.a, t1.b FROM (SELECT table1.* FROM table1 WHERE table1.e > CURRENT_DATE) AS t1 WHERE t1.a > 10"#,
    );
}
