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

fn main() {
    let sql = r#"SELECT current_timestamp()::TIMESTAMP_NTZ(9),
       payment_schedule_id,
       CASE
           WHEN records = 0 THEN 9
           ELSE 1
       END,
       CASE
           WHEN records = 0 THEN 'NO MEMBERS GENERATED FOR LIST'
           ELSE list_run_key
       END
FROM
  (SELECT lcl.payment_schedule_id,
          leb.list_run_key,
          count(distinct member_id) AS records
   FROM fraud.payment_list_schedule lcl
   LEFT JOIN fraud.payment_ui_list_export_batch leb ON lcl.payment_schedule_id = leb.payment_schedule_id
   GROUP BY 1,2) sub"#;

    let schema: HashMap<String, serde_json::Value> = serde_json::from_str(
        r#"{
        "PAYMENT_LIST_SCHEDULE": {"PAYMENT_SCHEDULE_ID": "VARCHAR"},
        "PAYMENT_UI_LIST_EXPORT_BATCH": {"LIST_RUN_KEY": "VARCHAR"}
    }"#,
    )
    .unwrap();

    // Raw edges
    let edges = polyglot_column_lineage(sql, SqlDialect::Snowflake).unwrap();
    eprintln!("=== Raw edges ({}) ===", edges.edges.len());
    for edge in &edges.edges {
        eprintln!(
            "  '{}'::'{}' -> '{}' ({:?})",
            edge.source_table, edge.source_column, edge.target_column, edge.terminal_kind
        );
    }

    // Processed
    let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();
    eprintln!("\n=== column_dict ===");
    for (k, v) in &result.column_dict {
        eprintln!("  {}: {:?}", k, v);
    }

    eprintln!("\n=== EXPECTED ===");
    eprintln!("  PAYMENT_SCHEDULE_ID: [PAYMENT_LIST_SCHEDULE.PAYMENT_SCHEDULE_ID]");
    eprintln!("  LIST_RUN_KEY: [PAYMENT_UI_LIST_EXPORT_BATCH.LIST_RUN_KEY]");
}
