use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use sqlparser::dialect::SnowflakeDialect;
use sqlparser::parser::Parser;
use std::collections::HashMap;

fn main() {
    // First, parse and inspect the AST
    let simple_sql = "SELECT y.value FROM t1, lateral flatten(input => col) y";
    let stmts = Parser::parse_sql(&SnowflakeDialect {}, simple_sql).unwrap();
    if let sqlparser::ast::Statement::Query(q) = &stmts[0] {
        if let sqlparser::ast::SetExpr::Select(sel) = q.body.as_ref() {
            for (i, twj) in sel.from.iter().enumerate() {
                eprintln!("from[{}] relation: {:#?}", i, twj.relation);
            }
        }
    }
    eprintln!("---");
    let sql = r#"select * from (
select TRACKING_TICKET_NUMBER,y.value::string as product_description
from (
select TRACKING_TICKET_NUMBER,parse_json(items) as items_j
from "ODS"."T1"
) X , lateral flatten(input => items_j) y
group by 1,2

union all

select tracking_ticket_number,product_description
from "ODS"."T2"
)"#;

    // Also check raw polyglot edges
    let poly_result =
        altimate_core::lineage::column_lineage::polyglot_column_lineage(sql, SqlDialect::Snowflake)
            .unwrap();
    eprintln!("Polyglot edges:");
    for e in &poly_result.edges {
        eprintln!(
            "  {} | {} -> {} (terminal: {:?})",
            e.source_table, e.source_column, e.target_column, e.terminal_kind
        );
    }
    eprintln!("---");

    let schema: HashMap<String, serde_json::Value> = serde_json::from_str(r#"{"T1": {"TRACKING_TICKET_NUMBER": "VARCHAR", "ITEMS": "VARCHAR"}, "T2": {"TRACKING_TICKET_NUMBER": "VARCHAR", "PRODUCT_DESCRIPTION": "VARCHAR"}}"#).unwrap();
    let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();
    let mut keys: Vec<_> = result.column_dict.keys().collect();
    keys.sort();
    for k in keys {
        let v = &result.column_dict[k];
        println!("  {} => {:?}", k, v);
    }
}
