//! Trace empty-result queries with CTEs to diagnose why column_dict is empty.
//!
//! Finds empty-actual queries from the ground truth fixture and prints a full
//! diagnostic trace showing:
//! 1. Raw polyglot-sql edges (source_table, source_column, target_column, terminal_kind)
//! 2. column_dict after projection_column_lineage
//! 3. CTE names extracted from the SQL
//! 4. Schema for the query
//! 5. Star key presence and sources
//!
//! Three traces:
//! - Case A: First empty-actual CTE+INSERT query (name contains "insert", SQL starts with WITH)
//! - Case B: First empty-actual CTE query where polyglot returns >0 edges (star expansion bug)
//! - Summary: Counts of empty-actual queries by root cause
//!
//! Run: cargo run --example trace_empty

use altimate_core::lineage::column_lineage;
use altimate_core::lineage::polyglot_column_lineage;
use altimate_core::schema::types::SqlDialect;
use serde::Deserialize;
use std::collections::HashMap;
use std::io::BufRead;
use std::path::Path;

use datafusion::sql::sqlparser::ast::{self, SetExpr, Statement};
use datafusion::sql::sqlparser::dialect::SnowflakeDialect;
use datafusion::sql::sqlparser::parser::Parser;

#[derive(Debug, Deserialize, Clone)]
struct Case {
    name: String,
    sql: String,
    schema: serde_json::Value,
    expected: HashMap<String, Vec<String>>,
}

fn schema_from_json(value: &serde_json::Value) -> HashMap<String, serde_json::Value> {
    match value {
        serde_json::Value::Object(map) => map.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
        _ => HashMap::new(),
    }
}

/// Extract CTE names from parsed statements.
fn extract_cte_names_from_stmts(stmts: &[Statement]) -> Vec<String> {
    let mut names = Vec::new();
    for stmt in stmts {
        let query = match stmt {
            Statement::Query(q) => q.as_ref(),
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    src.as_ref()
                } else {
                    continue;
                }
            }
            _ => continue,
        };
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                names.push(cte.alias.name.value.clone());
            }
        }
    }
    names
}

/// Extract CTE output column names from each CTE's SELECT list.
fn extract_cte_output_columns(stmts: &[Statement]) -> HashMap<String, Vec<String>> {
    let mut result = HashMap::new();
    for stmt in stmts {
        let query = match stmt {
            Statement::Query(q) => q.as_ref(),
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    src.as_ref()
                } else {
                    continue;
                }
            }
            _ => continue,
        };
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                let name = cte.alias.name.value.clone();
                let cols = select_column_names(&cte.query.body);
                result.insert(name, cols);
            }
        }
    }
    result
}

/// Extract column names from a SELECT list.
fn select_column_names(body: &SetExpr) -> Vec<String> {
    match body {
        SetExpr::Select(sel) => sel
            .projection
            .iter()
            .filter_map(|item| match item {
                ast::SelectItem::ExprWithAlias { alias, .. } => Some(alias.value.clone()),
                ast::SelectItem::UnnamedExpr(ast::Expr::Identifier(ident)) => {
                    Some(ident.value.clone())
                }
                ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(parts)) => {
                    parts.last().map(|p| p.value.clone())
                }
                ast::SelectItem::Wildcard(_) => Some("*".to_string()),
                ast::SelectItem::QualifiedWildcard(name, _) => Some(format!("{}.*", name)),
                _ => None,
            })
            .collect(),
        SetExpr::SetOperation { left, .. } => select_column_names(left),
        SetExpr::Query(q) => select_column_names(&q.body),
        _ => vec![],
    }
}

/// Extract FROM table references from the outermost SELECT.
fn extract_from_tables(stmts: &[Statement]) -> Vec<String> {
    let mut tables = Vec::new();
    for stmt in stmts {
        let query = match stmt {
            Statement::Query(q) => q.as_ref(),
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    src.as_ref()
                } else {
                    continue;
                }
            }
            _ => continue,
        };
        extract_from_tables_body(&query.body, &mut tables);
    }
    tables
}

fn extract_from_tables_body(body: &SetExpr, out: &mut Vec<String>) {
    if let SetExpr::Select(sel) = body {
        for item in &sel.from {
            extract_table_name(&item.relation, out);
            for join in &item.joins {
                extract_table_name(&join.relation, out);
            }
        }
    }
}

fn extract_table_name(factor: &ast::TableFactor, out: &mut Vec<String>) {
    match factor {
        ast::TableFactor::Table { name, alias, .. } => {
            let table_name = name.to_string();
            if let Some(a) = alias {
                out.push(format!("{} AS {}", table_name, a.name.value));
            } else {
                out.push(table_name);
            }
        }
        ast::TableFactor::Derived { alias, .. } => {
            if let Some(a) = alias {
                out.push(format!("(subquery) AS {}", a.name.value));
            } else {
                out.push("(subquery)".to_string());
            }
        }
        _ => {}
    }
}

/// Extract table names from a SetExpr, recursing into derived subqueries.
/// Matches the library's `extract_table_names_from_body` + `extract_table_name_from_factor`.
fn extract_tables_from_set_expr(body: &SetExpr) -> Vec<String> {
    match body {
        SetExpr::Select(sel) => {
            let mut tables = Vec::new();
            for item in &sel.from {
                extract_table_from_factor(&item.relation, &mut tables);
                for join in &item.joins {
                    extract_table_from_factor(&join.relation, &mut tables);
                }
            }
            tables
        }
        SetExpr::SetOperation { left, right, .. } => {
            let mut t = extract_tables_from_set_expr(left);
            t.extend(extract_tables_from_set_expr(right));
            t
        }
        SetExpr::Query(q) => extract_tables_from_set_expr(&q.body),
        _ => vec![],
    }
}

fn extract_table_from_factor(factor: &ast::TableFactor, out: &mut Vec<String>) {
    match factor {
        ast::TableFactor::Table { name, .. } => {
            out.push(name.to_string());
        }
        ast::TableFactor::Derived { subquery, .. } => {
            out.extend(extract_tables_from_set_expr(&subquery.body));
        }
        _ => {}
    }
}

/// Detect if the query body is a set operation and return left-side tables.
fn detect_set_operation_info(body: &SetExpr) -> (bool, Vec<String>, Vec<String>) {
    match body {
        SetExpr::SetOperation { left, right, .. } => {
            let left_tables = extract_tables_from_set_expr(left);
            let right_tables = extract_tables_from_set_expr(right);
            (true, left_tables, right_tables)
        }
        _ => (false, vec![], vec![]),
    }
}

/// Print a full diagnostic trace for a single case.
fn trace_case(case: &Case, dialect: SqlDialect, label: &str) {
    let schema = schema_from_json(&case.schema);
    let sf_dialect = SnowflakeDialect {};

    println!("================================================================");
    println!("  {}", label);
    println!("================================================================");
    println!();
    println!("Name: {}", case.name);
    println!();

    // SQL preview
    let sql_preview: String = case.sql.chars().take(600).collect();
    println!("--- SQL (first 600 chars) ---");
    println!("{}", sql_preview);
    if case.sql.len() > 600 {
        println!("... ({} total chars)", case.sql.len());
    }
    println!();

    // Schema
    println!("--- Schema ---");
    for (table, cols) in &schema {
        if let serde_json::Value::Object(obj) = cols {
            let col_names: Vec<&String> = obj.keys().collect();
            println!("  {}: {:?}", table, col_names);
        } else {
            println!("  {}: {:?}", table, cols);
        }
    }
    println!();

    // Expected
    println!("--- Expected ({} columns) ---", case.expected.len());
    let mut exp_sorted: Vec<_> = case.expected.iter().collect();
    exp_sorted.sort_by(|(a, _), (b, _)| a.cmp(b));
    for (col, sources) in &exp_sorted {
        println!("  {}: {:?}", col, sources);
    }
    println!();

    // Parse SQL for CTE info
    let stmts = Parser::parse_sql(&sf_dialect, &case.sql);
    match &stmts {
        Ok(s) => {
            let cte_names = extract_cte_names_from_stmts(s);
            println!("--- CTE names ({}) ---", cte_names.len());
            for name in &cte_names {
                println!("  {}", name);
            }
            println!();

            let cte_output = extract_cte_output_columns(s);
            println!("--- CTE output columns ---");
            for (cte, cols) in &cte_output {
                println!("  {}: {:?}", cte, cols);
            }
            println!();

            let outer_cols = match s.first() {
                Some(Statement::Query(q)) => select_column_names(&q.body),
                Some(Statement::Insert(ins)) => {
                    if let Some(ref src) = ins.source {
                        select_column_names(&src.body)
                    } else {
                        vec![]
                    }
                }
                _ => vec![],
            };
            println!("--- Outer SELECT columns ---");
            println!("  {:?}", outer_cols);
            println!();

            let from_tables = extract_from_tables(s);
            println!("--- Outer FROM tables ---");
            for t in &from_tables {
                println!("  {}", t);
            }
            println!();

            // Set operation detection
            let query_body = match s.first() {
                Some(Statement::Query(q)) => Some((&q.body, q.as_ref())),
                Some(Statement::Insert(ins)) => {
                    ins.source.as_ref().map(|src| (&src.body, src.as_ref()))
                }
                _ => None,
            };
            if let Some((body, query)) = query_body {
                let (is_set_op, left_tables, right_tables) = detect_set_operation_info(body);
                println!("--- Set operation detection ---");
                println!("  is_set_operation: {}", is_set_op);
                if is_set_op {
                    println!("  left_tables: {:?}", left_tables);
                    println!("  right_tables: {:?}", right_tables);
                }
                println!();

                // Simulate cte_source_tables (transitive resolution)
                if let Some(with) = &query.with {
                    let cte_name_set: std::collections::BTreeSet<String> = with
                        .cte_tables
                        .iter()
                        .map(|c| c.alias.name.value.clone())
                        .collect();

                    let mut cte_src_map: std::collections::BTreeMap<String, Vec<String>> =
                        std::collections::BTreeMap::new();
                    for cte in &with.cte_tables {
                        let name = cte.alias.name.value.clone();
                        let tables = extract_tables_from_set_expr(&cte.query.body);
                        cte_src_map.insert(name, tables);
                    }

                    // Transitive resolution (single pass)
                    let cte_order: Vec<String> = with
                        .cte_tables
                        .iter()
                        .map(|c| c.alias.name.value.clone())
                        .collect();
                    for cte_name in &cte_order {
                        let tables = cte_src_map.get(cte_name).cloned().unwrap_or_default();
                        let mut resolved = Vec::new();
                        for t in &tables {
                            if cte_name_set.contains(t) {
                                if let Some(inner) = cte_src_map.get(t) {
                                    resolved.extend(inner.clone());
                                }
                            } else {
                                resolved.push(t.clone());
                            }
                        }
                        resolved.sort();
                        resolved.dedup();
                        cte_src_map.insert(cte_name.clone(), resolved);
                    }

                    println!("--- CTE source tables (transitive resolution) ---");
                    for (cte, sources) in &cte_src_map {
                        println!("  {} -> {:?}", cte, sources);
                    }
                    println!();
                }
            }
        }
        Err(e) => {
            println!("Parse error: {}", e);
            println!();
        }
    }

    // Raw polyglot-sql edges
    println!("--- Raw polyglot-sql edges ---");
    match polyglot_column_lineage(&case.sql, dialect) {
        Ok(raw) => {
            println!("  Total edges: {}", raw.edges.len());
            println!("  Errors: {}", raw.errors.len());
            for err in &raw.errors {
                println!("    error[{}]: {}", err.column, err.message);
            }
            println!();

            for (i, edge) in raw.edges.iter().enumerate() {
                println!(
                    "  edge[{:>3}] src_table={:<30} src_col={:<25} tgt_col={:<25} kind={:?}",
                    i,
                    format!("\"{}\"", edge.source_table),
                    format!("\"{}\"", edge.source_column),
                    format!("\"{}\"", edge.target_column),
                    edge.terminal_kind
                );
            }

            // Group by target
            println!();
            println!("--- Edge summary by target column ---");
            let mut by_target: HashMap<String, Vec<usize>> = HashMap::new();
            for (i, edge) in raw.edges.iter().enumerate() {
                by_target
                    .entry(edge.target_column.clone())
                    .or_default()
                    .push(i);
            }
            let mut targets: Vec<_> = by_target.iter().collect();
            targets.sort_by(|(a, _), (b, _)| a.cmp(b));
            for (target, indices) in &targets {
                println!("  target \"{}\" <- {} edges", target, indices.len());
                for &i in *indices {
                    let e = &raw.edges[i];
                    println!(
                        "    [{:>3}] {}.{} ({:?})",
                        i, e.source_table, e.source_column, e.terminal_kind
                    );
                }
            }

            // Star check
            let has_star = raw.edges.iter().any(|e| e.target_column == "*");
            println!();
            println!("  Has '*' target edge: {}", has_star);
            if has_star {
                println!("  Star edges:");
                for edge in raw.edges.iter().filter(|e| e.target_column == "*") {
                    let in_schema = schema.contains_key(&edge.source_table)
                        || schema
                            .keys()
                            .any(|k| k.eq_ignore_ascii_case(&edge.source_table));
                    println!(
                        "    source_table=\"{}\" source_column=\"{}\" kind={:?} in_schema={}",
                        edge.source_table, edge.source_column, edge.terminal_kind, in_schema
                    );
                }
            }
        }
        Err(e) => {
            println!("  column_lineage error: {}", e);
        }
    }
    println!();

    // projection_column_lineage result
    let result = column_lineage(&case.sql, dialect, &schema, None);
    match &result {
        Ok(r) => {
            println!("--- projection_column_lineage result ---");
            println!("  column_dict: {} entries", r.column_dict.len());
            for (k, v) in &r.column_dict {
                println!("    {}: {:?}", k, v);
            }
            println!("  source_tables: {:?}", r.source_tables);
            println!("  errors: {:?}", r.errors);
        }
        Err(e) => {
            println!("  projection_column_lineage error: {}", e);
        }
    }
    println!();

    // Diagnosis
    println!("--- Diagnosis ---");
    let is_empty = result
        .as_ref()
        .map(|r| r.column_dict.is_empty())
        .unwrap_or(true);
    if is_empty {
        match polyglot_column_lineage(&case.sql, dialect) {
            Ok(raw) => {
                if raw.edges.is_empty() {
                    println!("  ROOT CAUSE: polyglot-sql returned ZERO edges.");
                    println!("  The query's lineage is not being traced by polyglot-sql.");
                    println!("  This is likely an INSERT INTO ... SELECT pattern that");
                    println!("  polyglot-sql does not support or a deeply nested CTE chain.");
                } else {
                    let star_count = raw.edges.iter().filter(|e| e.target_column == "*").count();
                    let table_count = raw
                        .edges
                        .iter()
                        .filter(|e| e.terminal_kind == polyglot_sql::lineage::TerminalKind::Table)
                        .count();
                    let unknown_count = raw
                        .edges
                        .iter()
                        .filter(|e| e.terminal_kind == polyglot_sql::lineage::TerminalKind::Unknown)
                        .count();

                    println!(
                        "  Polyglot returned {} edges ({} Table, {} Star, {} Unknown)",
                        raw.edges.len(),
                        table_count,
                        star_count,
                        unknown_count
                    );

                    // Check if all source tables are CTEs (not in schema)
                    let unique_sources: std::collections::BTreeSet<String> =
                        raw.edges.iter().map(|e| e.source_table.clone()).collect();
                    println!("  Unique source tables: {:?}", unique_sources);
                    for src in &unique_sources {
                        let in_schema = schema.contains_key(src)
                            || schema.keys().any(|k| k.eq_ignore_ascii_case(src));
                        println!("    \"{}\" -> in_schema={}", src, in_schema);
                    }

                    if star_count > 0 {
                        println!();
                        println!("  LIKELY CAUSE: Star edges present but expand_stars failed.");
                        println!();

                        // Check if this is a set operation with CTE star sources
                        if let Ok(stmts) = Parser::parse_sql(&sf_dialect, &case.sql) {
                            let query_body = match stmts.first() {
                                Some(Statement::Query(q)) => Some(q.as_ref()),
                                Some(Statement::Insert(ins)) => {
                                    ins.source.as_ref().map(|s| s.as_ref())
                                }
                                _ => None,
                            };
                            if let Some(q) = query_body {
                                let (is_set_op, left_tables, _right_tables) =
                                    detect_set_operation_info(&q.body);

                                if is_set_op {
                                    println!("  SET OPERATION DETECTED:");
                                    println!("    left_tables (from parser): {:?}", left_tables);
                                    println!();

                                    // Simulate expand_cte_star flow
                                    let cte_names = extract_cte_names_from_stmts(&stmts);
                                    let cte_output = extract_cte_output_columns(&stmts);

                                    for edge in raw.edges.iter().filter(|e| e.target_column == "*")
                                    {
                                        let src = &edge.source_table;
                                        let is_in_left = left_tables
                                            .iter()
                                            .any(|t| t.to_lowercase() == src.to_lowercase());
                                        println!(
                                            "    Star source \"{}\" -> in_left_tables={}",
                                            src, is_in_left
                                        );

                                        if !is_in_left {
                                            println!(
                                                "      FILTERED OUT by set-op left_tables filter"
                                            );
                                            continue;
                                        }

                                        // Check if it's a CTE
                                        let is_cte =
                                            cte_names.iter().any(|c| c.eq_ignore_ascii_case(src));
                                        println!("      is_cte={}", is_cte);

                                        if is_cte {
                                            // What are the CTE's output columns?
                                            let output = cte_output
                                                .iter()
                                                .find(|(k, _)| k.eq_ignore_ascii_case(src))
                                                .map(|(_, v)| v.clone())
                                                .unwrap_or_default();
                                            println!("      cte_output_columns={:?}", output);

                                            if output.contains(&"*".to_string()) {
                                                println!(
                                                    "      CTE has SELECT * -> expand through cte_source_tables"
                                                );
                                                println!(
                                                    "      BUG: expand_table_star on resolved base tables"
                                                );
                                                println!(
                                                    "      will fail the set-op left_tables filter because"
                                                );
                                                println!(
                                                    "      left_tables={:?} but base tables are schema-level",
                                                    left_tables
                                                );
                                                println!(
                                                    "      names that don't appear in left_tables."
                                                );
                                            }
                                        }
                                    }
                                } else {
                                    println!(
                                        "  Not a set operation. Star expansion bug is elsewhere."
                                    );
                                    println!("  Check CTE name resolution and schema matching.");
                                }
                            }
                        }
                    } else if table_count > 0 && unknown_count == 0 {
                        println!();
                        println!("  LIKELY CAUSE: All edges are Table kind, but CTE resolution");
                        println!("  could not map them to schema tables. The CTE source_table");
                        println!("  names may not be present in cte_source_tables map.");
                    } else if unknown_count > 0 {
                        println!();
                        println!("  LIKELY CAUSE: Unknown-kind edges not resolved. Source tables");
                        println!("  may be subquery aliases or unrecognized references.");
                    }
                }
            }
            Err(e) => {
                println!("  ROOT CAUSE: column_lineage failed: {}", e);
            }
        }
    }

    println!();
    println!("================================================================");
    println!();
}

fn main() {
    let fixture_path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../../tests/fixtures/lineage/snowflake_ground_truth.jsonl");

    let dialect = SqlDialect::Snowflake;

    // Collect all cases first (fixture is small enough)
    let file = std::fs::File::open(&fixture_path).expect("failed to open fixture");
    let reader = std::io::BufReader::with_capacity(1024 * 1024, file);
    let mut cases: Vec<Case> = Vec::new();
    for line in reader.lines() {
        let line = line.expect("failed to read line");
        let line = line.trim().to_string();
        if line.is_empty() {
            continue;
        }
        if let Ok(case) = serde_json::from_str::<Case>(&line) {
            cases.push(case);
        }
    }
    eprintln!("Loaded {} fixture cases.", cases.len());

    // ===== Case A: First empty-actual CTE+INSERT query =====
    let mut case_a: Option<Case> = None;
    let mut case_b: Option<Case> = None;

    // Summary counters
    let mut total_empty = 0u64;
    let mut empty_zero_edges = 0u64;
    let mut empty_has_edges_no_dict = 0u64;
    let mut empty_has_star_edges = 0u64;
    let mut empty_has_cte = 0u64;
    let mut empty_is_insert = 0u64;

    for case in &cases {
        if case.expected.is_empty() {
            continue;
        }

        let schema = schema_from_json(&case.schema);
        let result = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => r,
            Err(_) => continue,
        };

        if !result.column_dict.is_empty() {
            continue;
        }

        total_empty += 1;

        let sql_upper = case.sql.trim().to_uppercase();
        let has_cte = sql_upper.starts_with("WITH");
        let name_lower = case.name.to_lowercase();
        let is_insert = name_lower.contains("insert")
            || sql_upper.contains("INSERT INTO")
            || sql_upper.contains("INSERT OVERWRITE");

        if has_cte {
            empty_has_cte += 1;
        }
        if is_insert {
            empty_is_insert += 1;
        }

        // Check raw edges
        match polyglot_column_lineage(&case.sql, dialect) {
            Ok(raw) => {
                if raw.edges.is_empty() {
                    empty_zero_edges += 1;
                } else {
                    empty_has_edges_no_dict += 1;
                    let has_star = raw.edges.iter().any(|e| e.target_column == "*");
                    if has_star {
                        empty_has_star_edges += 1;
                    }

                    // Case B: CTE query with edges but empty column_dict
                    if case_b.is_none() && has_cte {
                        case_b = Some(case.clone());
                    }
                }
            }
            Err(_) => {
                empty_zero_edges += 1;
            }
        }

        // Case A: CTE+INSERT
        if case_a.is_none() && has_cte && is_insert {
            case_a = Some(case.clone());
        }
    }

    // Print summary first
    println!("################################################################");
    println!("  SUMMARY: Empty-actual queries");
    println!("################################################################");
    println!();
    println!("  Total empty-actual queries:          {}", total_empty);
    println!(
        "  Zero edges from polyglot:            {}",
        empty_zero_edges
    );
    println!(
        "  Has edges but empty column_dict:     {}",
        empty_has_edges_no_dict
    );
    println!(
        "    of which have star edges:           {}",
        empty_has_star_edges
    );
    println!("  Has CTE (WITH):                      {}", empty_has_cte);
    println!("  Is INSERT:                           {}", empty_is_insert);
    println!();

    // Trace Case A
    if let Some(ref case) = case_a {
        trace_case(case, dialect, "CASE A: First empty-actual CTE+INSERT query");
    } else {
        println!("No CTE+INSERT empty-actual query found.");
        println!();
    }

    // Trace Case B
    if let Some(ref case) = case_b {
        trace_case(
            case,
            dialect,
            "CASE B: First empty-actual CTE query with edges (star expansion bug?)",
        );
    } else {
        println!("No CTE query with edges but empty column_dict found.");
        println!();
    }
}
