//! SQL visualization DAG engine.
//!
//! Builds a directed acyclic graph representing the structural decomposition
//! of a SQL query. Each node represents a table, join, filter, group, or
//! subquery; edges represent data-flow dependencies.
//!
//! The output matches the Python `get_cte_lineage` / `get_sql_lineage_from_qualified`
//! format for frontend compatibility.

use super::types::{DagColumn, DagNode, DagResult, TableSchema};
use polyglot_sql::dialects::DialectType;
use polyglot_sql::expressions::{Expression, JoinKind};
use polyglot_sql::generator::Generator;
use regex::Regex;
use std::collections::{HashMap, VecDeque};

const FINAL_SELECT: &str = "__final_select__";
const DAG_THREAD_STACK_SIZE: usize = 64 * 1024 * 1024;

/// Build a SQL visualization DAG from a SQL string.
///
/// Parses the SQL using the given dialect, extracts CTEs and the main query,
/// and produces a graph of structural nodes (tables, joins, filters, etc.).
pub fn build_dag(sql: &str, dialect: DialectType) -> Result<DagResult, String> {
    build_dag_with_schema(sql, dialect, None)
}

/// Build a SQL visualization DAG with optional schema enrichment.
///
/// When `schema` is provided, table nodes are enriched with column names and types.
pub fn build_dag_with_schema(
    sql: &str,
    dialect: DialectType,
    schema: Option<&TableSchema>,
) -> Result<DagResult, String> {
    let sql_owned = sql.to_string();
    // Schema must be cloned into the thread since we can't send a reference.
    let schema_owned = schema.cloned();
    std::thread::Builder::new()
        .stack_size(DAG_THREAD_STACK_SIZE)
        .spawn(move || build_dag_inner(&sql_owned, dialect, schema_owned.as_ref()))
        .map_err(|e| format!("DAG thread spawn failed: {e}"))?
        .join()
        .map_err(|_| "DAG thread panicked".to_string())?
}

/// Convert `//` line comments to `--` so parsers don't choke.
fn strip_slash_comments(sql: &str) -> String {
    Regex::new(r"(?m)//([^\n]*)")
        .unwrap()
        .replace_all(sql, "-- $1")
        .to_string()
}

/// Sentinel value indicating a DML with no queryable body (e.g., INSERT VALUES).
const EMPTY_BODY: &str = "__empty_dml_body__";

/// Check if an expression is the empty-body sentinel.
fn is_empty_body(expr: &Expression) -> bool {
    matches!(expr, Expression::Literal(polyglot_sql::expressions::Literal::String(s)) if s == EMPTY_BODY)
}

/// Create the empty-body sentinel expression.
fn empty_body_expr() -> Expression {
    Expression::Literal(polyglot_sql::expressions::Literal::String(
        EMPTY_BODY.to_string(),
    ))
}

/// Unwrap a DML expression to extract (stmt_name, SELECT body).
///
/// Returns the target table name (uppercased) and the inner query expression.
/// For plain SELECT/UNION/etc, returns `("", expr)`.
fn unwrap_dml(expr: &Expression) -> (String, Expression) {
    match expr {
        Expression::Insert(ins) => {
            let stmt_name = table_ref_last_name(&ins.table);
            if let Some(ref query) = ins.query {
                (stmt_name, query.clone())
            } else {
                (stmt_name, empty_body_expr())
            }
        }
        Expression::CreateTable(ct) => {
            let stmt_name = table_ref_last_name(&ct.name);
            if let Some(ref sel) = ct.as_select {
                (stmt_name, sel.clone())
            } else {
                (stmt_name, empty_body_expr())
            }
        }
        Expression::Merge(m) => {
            let stmt_name = match &*m.this {
                Expression::Table(t) => table_ref_last_name(t),
                _ => "MERGE_TARGET".to_string(),
            };
            (stmt_name, (*m.using).clone())
        }
        Expression::Update(upd) => {
            let stmt_name = table_ref_last_name(&upd.table);
            if let Some(ref from) = upd.from_clause {
                (stmt_name, Expression::From(Box::new(from.clone())))
            } else {
                (stmt_name, Expression::Table(upd.table.clone()))
            }
        }
        Expression::Delete(del) => {
            let stmt_name = table_ref_last_name(&del.table);
            (stmt_name, Expression::Table(del.table.clone()))
        }
        _ => (String::new(), expr.clone()),
    }
}

/// Extract the last (unqualified) segment of a table name, uppercased.
fn table_ref_last_name(t: &polyglot_sql::expressions::TableRef) -> String {
    t.name.name.to_uppercase()
}

/// Try parsing with sqlparser as a fallback when polyglot-sql fails.
///
/// Extracts the DML target and SELECT body text, then re-parses the body
/// with polyglot-sql for dialect-aware AST traversal.
fn sqlparser_fallback(sql: &str, dialect: DialectType) -> Result<(String, Expression), String> {
    use sqlparser::dialect::GenericDialect;
    use sqlparser::parser::Parser;

    let statements = Parser::parse_sql(&GenericDialect {}, sql)
        .map_err(|e| format!("sqlparser fallback failed: {e}"))?;

    let stmt = statements
        .first()
        .ok_or_else(|| "sqlparser: no statements found".to_string())?;

    let (target_name, body_text) = match stmt {
        sqlparser::ast::Statement::Insert(ins) => {
            let target = extract_table_object_name(&ins.table);
            let body = ins
                .source
                .as_ref()
                .map(|q| q.to_string())
                .ok_or_else(|| "INSERT has no SELECT body".to_string())?;
            (target, body)
        }
        sqlparser::ast::Statement::CreateTable(ct) => {
            let target = extract_object_name_sqlparser(&ct.name);
            let body = ct
                .query
                .as_ref()
                .map(|q| q.to_string())
                .ok_or_else(|| "CREATE TABLE has no AS SELECT".to_string())?;
            (target, body)
        }
        sqlparser::ast::Statement::Merge { table, source, .. } => {
            let target = extract_table_factor_name(table);
            let body = source.to_string();
            (target, body)
        }
        _ => return Err("sqlparser: unsupported statement type".to_string()),
    };

    let stmt_name = target_name
        .split('.')
        .next_back()
        .unwrap_or(&target_name)
        .to_uppercase();

    // Re-parse the extracted body with polyglot-sql
    let body_expr = polyglot_sql::parse_one(&body_text, dialect)
        .map_err(|e| format!("polyglot-sql re-parse of body failed: {e}"))?;

    Ok((stmt_name, body_expr))
}

/// Extract table name from sqlparser `TableObject`.
fn extract_table_object_name(table: &sqlparser::ast::TableObject) -> String {
    match table {
        sqlparser::ast::TableObject::TableName(name) => extract_object_name_sqlparser(name),
        sqlparser::ast::TableObject::TableFunction(f) => f.to_string(),
    }
}

/// Extract table name from sqlparser `TableFactor`.
fn extract_table_factor_name(factor: &sqlparser::ast::TableFactor) -> String {
    match factor {
        sqlparser::ast::TableFactor::Table { name, .. } => extract_object_name_sqlparser(name),
        other => other.to_string(),
    }
}

/// Extract object name from sqlparser ObjectName (e.g., `db.schema.table`).
fn extract_object_name_sqlparser(name: &sqlparser::ast::ObjectName) -> String {
    name.0
        .iter()
        .map(|part| match part {
            sqlparser::ast::ObjectNamePart::Identifier(ident) => ident.value.clone(),
        })
        .collect::<Vec<_>>()
        .join(".")
}

/// Enrich table nodes with column metadata from schema.
fn enrich_with_schema(result: &mut DagResult, schema: &TableSchema) {
    for node in result.nodes.values_mut() {
        if node.kind != "table" {
            continue;
        }
        let name = match node.name {
            Some(ref n) => n.clone(),
            None => continue,
        };

        let db = node.db.as_deref().unwrap_or("");
        let sch = node.schema.as_deref().unwrap_or("");

        // Try exact lookup with decreasing specificity
        let mut keys: Vec<String> = Vec::new();
        if !db.is_empty() && !sch.is_empty() {
            keys.push(format!("{db}.{sch}.{name}").to_uppercase());
        }
        if !sch.is_empty() {
            keys.push(format!("{sch}.{name}").to_uppercase());
        }
        keys.push(name.to_uppercase());

        let mut matched = false;
        for key in &keys {
            if let Some(table_cols) = schema.get(key) {
                node.columns = cols_from_schema(table_cols);
                matched = true;
                break;
            }
        }

        // Fallback: suffix-match against schema keys (handles SQL with 2-part
        // names like `fraud.table` when schema keys are 3-part `DB.FRAUD.TABLE`)
        if !matched {
            let suffix2 = if !sch.is_empty() {
                Some(format!(".{sch}.{name}").to_uppercase())
            } else {
                None
            };
            let suffix1 = format!(".{}", name.to_uppercase());

            for (schema_key, table_cols) in schema.iter() {
                let sk = schema_key.to_uppercase();
                let hit = suffix2.as_ref().is_some_and(|s| sk.ends_with(s.as_str()))
                    || sk.ends_with(&suffix1);
                if hit {
                    node.columns = cols_from_schema(table_cols);
                    break;
                }
            }
        }
    }
}

fn cols_from_schema(table_cols: &HashMap<String, String>) -> Vec<DagColumn> {
    table_cols
        .iter()
        .map(|(col_name, dtype)| DagColumn {
            name: col_name.clone(),
            expression: None,
            data_type: Some(dtype.clone()),
        })
        .collect()
}

fn build_dag_inner(
    sql: &str,
    dialect: DialectType,
    schema: Option<&TableSchema>,
) -> Result<DagResult, String> {
    let sql = strip_slash_comments(sql);

    // Try polyglot-sql first; fall back to sqlparser on failure
    let (stmt_name, body, cte_source) = match polyglot_sql::parse_one(&sql, dialect) {
        Ok(expr) => {
            let (name, body) = unwrap_dml(&expr);
            (name, body, expr)
        }
        Err(polyglot_err) => match sqlparser_fallback(&sql, dialect) {
            Ok((name, body_expr)) => (name, body_expr.clone(), body_expr),
            Err(fallback_err) => {
                return Err(format!(
                    "parse error: {polyglot_err} (fallback: {fallback_err})"
                ));
            }
        },
    };

    let has_empty_body = is_empty_body(&body);

    let mut all_nodes: HashMap<String, DagNode> = HashMap::new();
    let mut all_deps: HashMap<String, Vec<String>> = HashMap::new();
    let mut ctes = Vec::new();

    if !has_empty_body {
        // Extract CTEs from the original expression (before DML unwrap)
        extract_ctes(
            &cte_source,
            &mut all_nodes,
            &mut all_deps,
            &mut ctes,
            dialect,
        );
        // Also check for CTEs on the body itself (CTAS with CTE in the AS SELECT)
        if !std::ptr::eq(&cte_source as *const _, &body as *const _) {
            extract_ctes(&body, &mut all_nodes, &mut all_deps, &mut ctes, dialect);
        }

        // Build DAG for the main query body (with CTEs stripped)
        let main_body = strip_with(&body);
        let mut builder = DagBuilder::new(FINAL_SELECT, ctes.clone());
        builder.recurse(&main_body, FINAL_SELECT, dialect);
        merge_into(&mut all_nodes, &mut all_deps, builder);

        // Add final select node metadata
        let main_sql = gen_pretty(&main_body, dialect);
        all_nodes.insert(
            FINAL_SELECT.to_string(),
            DagNode {
                node_id: FINAL_SELECT.to_string(),
                kind: "final".to_string(),
                id: Some("0".to_string()),
                name: Some(FINAL_SELECT.to_string()),
                columns: extract_select_columns(&main_body, dialect),
                expression: main_sql.clone(),
                sql: Some(main_sql),
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
    } else {
        // Minimal DAG for DML without a queryable body (e.g., INSERT VALUES, plain CREATE TABLE)
        all_deps.insert(FINAL_SELECT.to_string(), Vec::new());
        all_nodes.insert(
            FINAL_SELECT.to_string(),
            DagNode {
                node_id: FINAL_SELECT.to_string(),
                kind: "final".to_string(),
                id: Some("0".to_string()),
                name: Some(FINAL_SELECT.to_string()),
                columns: Vec::new(),
                expression: String::new(),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
    }

    // Assign hierarchical BFS IDs
    assign_bfs_ids(&mut all_nodes, &all_deps);

    let mut result = DagResult {
        stmt_name,
        dependencies: all_deps,
        nodes: all_nodes,
        ctes,
    };

    // Schema enrichment (post-process)
    if let Some(schema) = schema {
        enrich_with_schema(&mut result, schema);
    }

    Ok(result)
}

// ─── CTE extraction ─────────────────────────────────────────────

/// Walk the expression tree to find and process CTEs.
fn extract_ctes(
    expr: &Expression,
    all_nodes: &mut HashMap<String, DagNode>,
    all_deps: &mut HashMap<String, Vec<String>>,
    ctes: &mut Vec<String>,
    dialect: DialectType,
) {
    let with = match expr {
        Expression::Select(sel) => sel.with.as_ref(),
        Expression::Union(u) => u.with.as_ref(),
        Expression::Intersect(i) => i.with.as_ref(),
        Expression::Except(e) => e.with.as_ref(),
        Expression::Insert(ins) => ins.query.as_ref().and_then(|q| extract_with(q)),
        Expression::CreateTable(ct) => ct
            .with_cte
            .as_ref()
            .or_else(|| ct.as_select.as_ref().and_then(|s| extract_with(s))),
        Expression::Update(upd) => upd.with.as_ref(),
        _ => None,
    };

    if let Some(w) = with {
        // Collect all CTE names first so each CTE can reference siblings
        let all_cte_names: Vec<String> = w.ctes.iter().map(|c| c.alias.name.clone()).collect();

        for cte in &w.ctes {
            let cte_name = cte.alias.name.clone();
            ctes.push(cte_name.clone());

            // Build DAG for this CTE body — pass sibling CTE names
            // so cross-CTE references resolve correctly
            let mut builder = DagBuilder::new(&cte_name, all_cte_names.clone());
            builder.recurse(&cte.this, &cte_name, dialect);
            merge_into(all_nodes, all_deps, builder);

            // Add CTE node metadata
            let cte_sql = gen_pretty(&cte.this, dialect);
            let columns = extract_select_columns(&cte.this, dialect);
            all_nodes.insert(
                cte_name.clone(),
                DagNode {
                    node_id: cte_name.clone(),
                    kind: "cte".to_string(),
                    id: None,
                    name: Some(cte_name),
                    columns,
                    expression: cte_sql.clone(),
                    sql: Some(cte_sql),
                    on_clause: None,
                    join_type: None,
                    db: None,
                    schema: None,
                },
            );
        }
    }
}

/// Extract the `With` clause from an expression, if any.
fn extract_with(expr: &Expression) -> Option<&polyglot_sql::expressions::With> {
    match expr {
        Expression::Select(sel) => sel.with.as_ref(),
        Expression::Union(u) => u.with.as_ref(),
        Expression::Intersect(i) => i.with.as_ref(),
        Expression::Except(e) => e.with.as_ref(),
        _ => None,
    }
}

/// Strip the WITH clause from a top-level expression, returning just the body.
fn strip_with(expr: &Expression) -> Expression {
    match expr {
        Expression::Select(sel) if sel.with.is_some() => {
            let mut sel = sel.as_ref().clone();
            sel.with = None;
            Expression::Select(Box::new(sel))
        }
        Expression::Union(u) if u.with.is_some() => {
            let mut u = u.as_ref().clone();
            u.with = None;
            Expression::Union(Box::new(u))
        }
        Expression::Intersect(i) if i.with.is_some() => {
            let mut i = i.as_ref().clone();
            i.with = None;
            Expression::Intersect(Box::new(i))
        }
        Expression::Except(e) if e.with.is_some() => {
            let mut e = e.as_ref().clone();
            e.with = None;
            Expression::Except(Box::new(e))
        }
        _ => expr.clone(),
    }
}

// ─── DAG builder ─────────────────────────────────────────────────

struct DagBuilder {
    nodes: HashMap<String, DagNode>,
    graph: HashMap<String, Vec<String>>,
    /// Per-prefix counters for deterministic 0-based node IDs (e.g., `Join_0`, `Join_1`).
    counters: HashMap<String, usize>,
    /// Known CTE names — bare table references matching these are treated as CTE refs.
    cte_names: Vec<String>,
}

impl DagBuilder {
    fn new(root_name: &str, cte_names: Vec<String>) -> Self {
        Self {
            nodes: HashMap::new(),
            graph: HashMap::from([(root_name.to_string(), Vec::new())]),
            counters: HashMap::new(),
            cte_names,
        }
    }

    fn next_id(&mut self, prefix: &str) -> String {
        let count = self.counters.entry(prefix.to_string()).or_insert(0);
        let id = format!("{prefix}_{count}");
        *count += 1;
        id
    }

    fn add_node(&mut self, parent_id: &str, node: DagNode) -> String {
        let node_id = node.node_id.clone();
        self.graph
            .entry(parent_id.to_string())
            .or_default()
            .push(node_id.clone());
        self.graph.entry(node_id.clone()).or_default();
        self.nodes.insert(node_id.clone(), node);
        node_id
    }

    /// Main recursive traversal — mirrors Python's `generate_dag.recurse()`.
    fn recurse(&mut self, expr: &Expression, parent_id: &str, dialect: DialectType) -> String {
        match expr {
            Expression::Table(t) => self.handle_table(t, parent_id, dialect),
            Expression::Select(s) => self.handle_select(s, parent_id, dialect),
            Expression::Union(u) => self.handle_union(u, parent_id, dialect),
            Expression::Intersect(i) => self.handle_intersect(i, parent_id, dialect),
            Expression::Except(e) => self.handle_except(e, parent_id, dialect),
            Expression::Join(j) => self.handle_join(j, parent_id, dialect),
            Expression::Where(w) => self.handle_where(w, parent_id, dialect),
            Expression::GroupBy(g) => self.handle_group(g, parent_id, dialect),
            Expression::Subquery(s) => self.handle_subquery(s, parent_id, dialect),
            Expression::From(f) => self.handle_from(f, parent_id, dialect),
            Expression::OrderBy(o) => self.handle_order(o, parent_id, dialect),
            _ => self.recurse_children(expr, parent_id, dialect),
        }
    }

    fn handle_table(
        &mut self,
        t: &polyglot_sql::expressions::TableRef,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let name = &t.name.name;
        if name.is_empty() {
            return parent_id.to_string();
        }

        let schema_name = t.schema.as_ref().map(|s| s.name.clone());
        let catalog_name = t.catalog.as_ref().map(|c| c.name.clone());

        // Build a qualified key for the node ID
        let key = if name.starts_with('@') {
            name.clone()
        } else if let Some(ref cat) = catalog_name {
            format!("{}.{}.{}", cat, schema_name.as_deref().unwrap_or(""), name)
        } else if let Some(ref sch) = schema_name {
            format!("{sch}.{name}")
        } else {
            name.clone()
        };

        // If no schema/catalog qualifier, check if this is a known CTE.
        // Store a reference to the CTE name as the dependency target
        // so the dependency graph links to the CTE node.
        if schema_name.is_none()
            && catalog_name.is_none()
            && !name.starts_with('@')
            && self.cte_names.iter().any(|c| c == name)
        {
            // CTE reference — link to existing CTE node
            self.graph
                .entry(parent_id.to_string())
                .or_default()
                .push(name.clone());
            return name.clone();
        }

        let node_id = self.next_id(&key);
        let sql_text = gen_sql(&Expression::Table(t.clone()), dialect);

        self.add_node(
            parent_id,
            DagNode {
                node_id: node_id.clone(),
                kind: "table".to_string(),
                id: None,
                name: Some(name.clone()),
                columns: Vec::new(),
                expression: sql_text,
                sql: None,
                on_clause: None,
                join_type: None,
                db: catalog_name,
                schema: schema_name,
            },
        )
    }

    fn handle_select(
        &mut self,
        sel: &polyglot_sql::expressions::Select,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let mut current_parent = parent_id.to_string();

        // Check if FROM clause exists
        let has_from = sel.from.is_some();
        if !has_from {
            // Standalone SELECT (e.g., SELECT 1, SELECT CURRENT_DATE)
            let node_id = self.next_id("Select");
            let columns = sel
                .expressions
                .iter()
                .map(|e| DagColumn {
                    name: extract_alias_or_name(e),
                    expression: Some(gen_sql(e, dialect)),
                    data_type: None,
                })
                .collect();
            return self.add_node(
                parent_id,
                DagNode {
                    node_id,
                    kind: "SELECT".to_string(),
                    id: None,
                    name: None,
                    columns,
                    expression: gen_sql(&Expression::Select(Box::new(sel.clone())), dialect),
                    sql: None,
                    on_clause: None,
                    join_type: None,
                    db: None,
                    schema: None,
                },
            );
        }

        // Decompose SELECT into a chain: GROUP BY → WHERE → JOINs → FROM
        // (reversed join order matches Python behavior)

        // 1. GROUP BY
        if let Some(ref group) = sel.group_by {
            let gid = self.next_id("Group");
            let group_expr = Expression::GroupBy(Box::new(group.clone()));
            current_parent = self.add_node(
                &current_parent,
                DagNode {
                    node_id: gid,
                    kind: "GROUP".to_string(),
                    id: None,
                    name: None,
                    columns: Vec::new(),
                    expression: gen_sql(&group_expr, dialect),
                    sql: None,
                    on_clause: None,
                    join_type: None,
                    db: None,
                    schema: None,
                },
            );
        }

        // 2. WHERE
        if let Some(ref wc) = sel.where_clause {
            let wid = self.next_id("Where");
            let where_expr = Expression::Where(Box::new(wc.clone()));
            let where_node_id = self.add_node(
                &current_parent,
                DagNode {
                    node_id: wid,
                    kind: "FILTER".to_string(),
                    id: None,
                    name: None,
                    columns: Vec::new(),
                    expression: gen_sql(&where_expr, dialect),
                    sql: None,
                    on_clause: None,
                    join_type: None,
                    db: None,
                    schema: None,
                },
            );
            // Recurse into WHERE to find table references in subqueries
            self.recurse_children(
                &Expression::Where(Box::new(wc.clone())),
                &where_node_id,
                dialect,
            );
            current_parent = where_node_id;
        }

        // 3. JOINs (reversed order per Python behavior)
        let joins_reversed: Vec<_> = sel.joins.iter().rev().collect();
        for join in joins_reversed {
            let join_id = self.handle_join(join, &current_parent, dialect);
            current_parent = join_id;
        }

        // 4. FROM
        if let Some(ref from) = sel.from {
            for from_expr in &from.expressions {
                let from_id = self.recurse(from_expr, &current_parent, dialect);
                current_parent = from_id;
            }
        }

        current_parent
    }

    fn handle_union(
        &mut self,
        u: &polyglot_sql::expressions::Union,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let node_id = self.next_id("Union");
        let columns = extract_select_columns_from_expr(&u.left, dialect);
        let union_node_id = self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "UNION".to_string(),
                id: None,
                name: None,
                columns,
                expression: gen_sql(&Expression::Union(Box::new(u.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
        self.recurse(&u.left, &union_node_id, dialect);
        self.recurse(&u.right, &union_node_id, dialect);
        union_node_id
    }

    fn handle_intersect(
        &mut self,
        i: &polyglot_sql::expressions::Intersect,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let node_id = self.next_id("Intersect");
        let intersect_node_id = self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "INTERSECT".to_string(),
                id: None,
                name: None,
                columns: Vec::new(),
                expression: gen_sql(&Expression::Intersect(Box::new(i.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
        self.recurse(&i.left, &intersect_node_id, dialect);
        self.recurse(&i.right, &intersect_node_id, dialect);
        intersect_node_id
    }

    fn handle_except(
        &mut self,
        e: &polyglot_sql::expressions::Except,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let node_id = self.next_id("Except");
        let except_node_id = self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "EXCEPT".to_string(),
                id: None,
                name: None,
                columns: Vec::new(),
                expression: gen_sql(&Expression::Except(Box::new(e.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
        self.recurse(&e.left, &except_node_id, dialect);
        self.recurse(&e.right, &except_node_id, dialect);
        except_node_id
    }

    fn handle_join(
        &mut self,
        j: &polyglot_sql::expressions::Join,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let kind_str = join_kind_label(j.kind);
        let on_text = j.on.as_ref().map(|on| gen_sql(on, dialect));

        let node_id = self.next_id("Join");
        let join_node_id = self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: format!("{kind_str}_JOIN"),
                id: None,
                name: None,
                columns: Vec::new(),
                expression: gen_sql(&Expression::Join(Box::new(j.clone())), dialect),
                sql: None,
                on_clause: on_text,
                join_type: Some(kind_str),
                db: None,
                schema: None,
            },
        );

        // Recurse into the joined table
        self.recurse(&j.this, &join_node_id, dialect);
        join_node_id
    }

    fn handle_where(
        &mut self,
        w: &polyglot_sql::expressions::Where,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let node_id = self.next_id("Where");
        let where_node_id = self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "FILTER".to_string(),
                id: None,
                name: None,
                columns: Vec::new(),
                expression: gen_sql(&Expression::Where(Box::new(w.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
        self.recurse_children(
            &Expression::Where(Box::new(w.clone())),
            &where_node_id,
            dialect,
        );
        where_node_id
    }

    fn handle_group(
        &mut self,
        g: &polyglot_sql::expressions::GroupBy,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let node_id = self.next_id("Group");
        self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "GROUP".to_string(),
                id: None,
                name: None,
                columns: Vec::new(),
                expression: gen_sql(&Expression::GroupBy(Box::new(g.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        )
    }

    fn handle_order(
        &mut self,
        o: &polyglot_sql::expressions::OrderBy,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let node_id = self.next_id("Sort");
        self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "SORT".to_string(),
                id: None,
                name: None,
                columns: Vec::new(),
                expression: gen_sql(&Expression::OrderBy(Box::new(o.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        )
    }

    fn handle_subquery(
        &mut self,
        s: &polyglot_sql::expressions::Subquery,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let alias = s.alias.as_ref().map(|a| a.name.clone());
        let name = alias.unwrap_or_else(|| "subquery".to_string());

        let node_id = self.next_id("Subquery");
        let sub_node_id = self.add_node(
            parent_id,
            DagNode {
                node_id,
                kind: "subquery".to_string(),
                id: None,
                name: Some(name),
                columns: Vec::new(),
                expression: gen_sql(&Expression::Subquery(Box::new(s.clone())), dialect),
                sql: None,
                on_clause: None,
                join_type: None,
                db: None,
                schema: None,
            },
        );
        self.recurse(&s.this, &sub_node_id, dialect);
        sub_node_id
    }

    fn handle_from(
        &mut self,
        f: &polyglot_sql::expressions::From,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        let mut last = parent_id.to_string();
        for expr in &f.expressions {
            last = self.recurse(expr, parent_id, dialect);
        }
        last
    }

    /// Recurse through child expressions without creating a node for the current level.
    fn recurse_children(
        &mut self,
        expr: &Expression,
        parent_id: &str,
        dialect: DialectType,
    ) -> String {
        use polyglot_sql::traversal::ExpressionWalk;
        for child in expr.children() {
            match child {
                Expression::Table(_)
                | Expression::Select(_)
                | Expression::Union(_)
                | Expression::Intersect(_)
                | Expression::Except(_)
                | Expression::Subquery(_) => {
                    self.recurse(child, parent_id, dialect);
                }
                _ => {
                    self.recurse_children(child, parent_id, dialect);
                }
            }
        }
        parent_id.to_string()
    }
}

// ─── BFS ID assignment ──────────────────────────────────────────

/// Assign hierarchical BFS IDs starting from `__final_select__`.
///
/// IDs are strings like `"0"`, `"00"`, `"01"`, `"010"`, etc., encoding
/// the BFS traversal position for frontend layout.
fn assign_bfs_ids(nodes: &mut HashMap<String, DagNode>, deps: &HashMap<String, Vec<String>>) {
    let mut visited: HashMap<String, bool> = HashMap::new();
    let mut queue: VecDeque<(String, String)> = VecDeque::new();

    // Start from final select
    if let Some(node) = nodes.get_mut(FINAL_SELECT) {
        node.id = Some("0".to_string());
    }
    queue.push_back(("0".to_string(), FINAL_SELECT.to_string()));

    while let Some((pos, current)) = queue.pop_front() {
        if visited.get(&current).copied().unwrap_or(false) {
            continue;
        }
        visited.insert(current.clone(), true);

        if let Some(children) = deps.get(&current) {
            for (i, child) in children.iter().enumerate() {
                let child_id = format!("{pos}{i}");
                if let Some(node) = nodes.get_mut(child) {
                    node.id = Some(child_id.clone());
                }
                queue.push_back((child_id, child.clone()));
            }
        }
    }
}

// ─── Helpers ────────────────────────────────────────────────────

fn merge_into(
    all_nodes: &mut HashMap<String, DagNode>,
    all_deps: &mut HashMap<String, Vec<String>>,
    builder: DagBuilder,
) {
    all_nodes.extend(builder.nodes);
    for (k, v) in builder.graph {
        all_deps.entry(k).or_insert(v);
    }
}

fn gen_sql(expr: &Expression, _dialect: DialectType) -> String {
    Generator::sql(expr).unwrap_or_default()
}

fn gen_pretty(expr: &Expression, _dialect: DialectType) -> String {
    Generator::pretty_sql(expr).unwrap_or_default()
}

fn join_kind_label(kind: JoinKind) -> String {
    match kind {
        JoinKind::Inner => "INNER".to_string(),
        JoinKind::Left => "LEFT".to_string(),
        JoinKind::Right => "RIGHT".to_string(),
        JoinKind::Full => "FULL".to_string(),
        JoinKind::Outer => "OUTER".to_string(),
        JoinKind::Cross => "CROSS".to_string(),
        JoinKind::Natural => "NATURAL".to_string(),
        JoinKind::NaturalLeft => "NATURAL_LEFT".to_string(),
        JoinKind::NaturalRight => "NATURAL_RIGHT".to_string(),
        JoinKind::NaturalFull => "NATURAL_FULL".to_string(),
        JoinKind::Semi => "SEMI".to_string(),
        JoinKind::Anti => "ANTI".to_string(),
        JoinKind::LeftSemi => "LEFT_SEMI".to_string(),
        JoinKind::LeftAnti => "LEFT_ANTI".to_string(),
        JoinKind::RightSemi => "RIGHT_SEMI".to_string(),
        JoinKind::RightAnti => "RIGHT_ANTI".to_string(),
        JoinKind::CrossApply => "CROSS_APPLY".to_string(),
        JoinKind::OuterApply => "OUTER_APPLY".to_string(),
        JoinKind::AsOf => "ASOF".to_string(),
        JoinKind::AsOfLeft => "ASOF_LEFT".to_string(),
        JoinKind::AsOfRight => "ASOF_RIGHT".to_string(),
        JoinKind::Lateral => "LATERAL".to_string(),
        JoinKind::LeftLateral => "LEFT_LATERAL".to_string(),
        JoinKind::Straight => "STRAIGHT".to_string(),
        JoinKind::Implicit => "CROSS".to_string(),
        JoinKind::Array => "ARRAY".to_string(),
        JoinKind::LeftArray => "LEFT_ARRAY".to_string(),
    }
}

/// Extract column names from a SELECT expression's select-list.
fn extract_select_columns(expr: &Expression, dialect: DialectType) -> Vec<DagColumn> {
    match expr {
        Expression::Select(sel) => sel
            .expressions
            .iter()
            .map(|e| DagColumn {
                name: extract_alias_or_name(e),
                expression: Some(gen_sql(e, dialect)),
                data_type: None,
            })
            .collect(),
        Expression::Union(u) => extract_select_columns(&u.left, dialect),
        Expression::Intersect(i) => extract_select_columns(&i.left, dialect),
        Expression::Except(e) => extract_select_columns(&e.left, dialect),
        _ => Vec::new(),
    }
}

fn extract_select_columns_from_expr(expr: &Expression, dialect: DialectType) -> Vec<DagColumn> {
    extract_select_columns(expr, dialect)
}

/// Extract alias or name from a select-list expression.
fn extract_alias_or_name(expr: &Expression) -> String {
    match expr {
        Expression::Alias(a) => a.alias.name.clone(),
        Expression::Column(c) => c.name.name.clone(),
        Expression::Star(_) => "*".to_string(),
        _ => gen_sql(expr, DialectType::Snowflake),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn with_large_stack<F: FnOnce() + Send + 'static>(f: F) {
        std::thread::Builder::new()
            .stack_size(64 * 1024 * 1024)
            .spawn(f)
            .unwrap()
            .join()
            .unwrap();
    }

    #[test]
    fn simple_select_produces_dag() {
        with_large_stack(|| {
            let result = build_dag("SELECT id, name FROM users", DialectType::Snowflake).unwrap();
            assert!(result.dependencies.contains_key(FINAL_SELECT));
            assert!(result.nodes.contains_key(FINAL_SELECT));
            assert!(result.ctes.is_empty());
            // Should have a table node for users
            let has_users = result
                .nodes
                .values()
                .any(|n| n.kind == "table" && n.name.as_deref() == Some("users"));
            assert!(has_users, "Should find users table node");
        });
    }

    #[test]
    fn join_produces_join_node() {
        with_large_stack(|| {
            let result = build_dag(
                "SELECT u.name, o.total FROM users u JOIN orders o ON u.id = o.user_id",
                DialectType::Snowflake,
            )
            .unwrap();
            let has_join = result.nodes.values().any(|n| n.kind.contains("JOIN"));
            assert!(has_join, "Should find a JOIN node");
        });
    }

    #[test]
    fn where_produces_filter_node() {
        with_large_stack(|| {
            let result = build_dag(
                "SELECT id FROM users WHERE age > 18",
                DialectType::Snowflake,
            )
            .unwrap();
            let has_filter = result.nodes.values().any(|n| n.kind == "FILTER");
            assert!(has_filter, "Should find a FILTER node");
        });
    }

    #[test]
    fn cte_produces_cte_node() {
        with_large_stack(|| {
            let result = build_dag(
                "WITH active AS (SELECT id FROM users WHERE status = 'active') \
                 SELECT * FROM active",
                DialectType::Snowflake,
            )
            .unwrap();
            assert!(result.ctes.contains(&"active".to_string()));
            assert!(result.nodes.contains_key("active"));
            let cte_node = &result.nodes["active"];
            assert_eq!(cte_node.kind, "cte");
        });
    }

    #[test]
    fn union_produces_union_node() {
        with_large_stack(|| {
            let result = build_dag(
                "SELECT id FROM users UNION ALL SELECT id FROM admins",
                DialectType::Snowflake,
            )
            .unwrap();
            let has_union = result.nodes.values().any(|n| n.kind == "UNION");
            assert!(has_union, "Should find a UNION node");
        });
    }

    #[test]
    fn group_by_produces_group_node() {
        with_large_stack(|| {
            let result = build_dag(
                "SELECT status, COUNT(*) FROM users GROUP BY status",
                DialectType::Snowflake,
            )
            .unwrap();
            let has_group = result.nodes.values().any(|n| n.kind == "GROUP");
            assert!(has_group, "Should find a GROUP node");
        });
    }

    #[test]
    fn subquery_produces_subquery_node() {
        with_large_stack(|| {
            let result = build_dag(
                "SELECT * FROM (SELECT id, name FROM users) sub",
                DialectType::Snowflake,
            )
            .unwrap();
            let has_sub = result.nodes.values().any(|n| n.kind == "subquery");
            assert!(has_sub, "Should find a subquery node");
        });
    }

    #[test]
    fn bfs_ids_assigned() {
        with_large_stack(|| {
            let result = build_dag(
                "SELECT id FROM users WHERE age > 18",
                DialectType::Snowflake,
            )
            .unwrap();
            // Final select should have id "0"
            assert_eq!(result.nodes[FINAL_SELECT].id.as_deref(), Some("0"));
            // All reachable nodes should have ids assigned
            for node in result.nodes.values() {
                assert!(
                    node.id.is_some(),
                    "Node {} should have BFS id",
                    node.node_id
                );
            }
        });
    }

    #[test]
    fn complex_cte_with_join() {
        with_large_stack(|| {
            let sql = "\
                WITH user_orders AS (\
                    SELECT u.id, u.name, o.total \
                    FROM users u \
                    JOIN orders o ON u.id = o.user_id \
                    WHERE o.status = 'complete'\
                ) \
                SELECT name, SUM(total) \
                FROM user_orders \
                GROUP BY name";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert!(result.ctes.contains(&"user_orders".to_string()));
            assert!(result.nodes.contains_key(FINAL_SELECT));
            // CTE should contain join and filter nodes
            let has_join = result.nodes.values().any(|n| n.kind.contains("JOIN"));
            let has_filter = result.nodes.values().any(|n| n.kind == "FILTER");
            assert!(has_join, "CTE should have a JOIN node");
            assert!(has_filter, "CTE should have a FILTER node");
        });
    }

    #[test]
    fn empty_sql_returns_error() {
        let result = build_dag("", DialectType::Snowflake);
        assert!(result.is_err());
    }

    #[test]
    fn final_select_has_columns() {
        with_large_stack(|| {
            let result =
                build_dag("SELECT id, name, email FROM users", DialectType::Snowflake).unwrap();
            let fs = &result.nodes[FINAL_SELECT];
            let col_names: Vec<&str> = fs.columns.iter().map(|c| c.name.as_str()).collect();
            assert!(col_names.contains(&"id"));
            assert!(col_names.contains(&"name"));
            assert!(col_names.contains(&"email"));
        });
    }

    // ─── DML tests ──────────────────────────────────────────────────

    #[test]
    fn insert_into_select_produces_dag() {
        with_large_stack(|| {
            let sql = "INSERT INTO target_table SELECT id, name FROM source_table";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert_eq!(result.stmt_name, "TARGET_TABLE");
            assert!(result.nodes.contains_key(FINAL_SELECT));
            let has_source = result
                .nodes
                .values()
                .any(|n| n.kind == "table" && n.name.as_deref() == Some("source_table"));
            assert!(has_source, "Should find source_table node");
        });
    }

    #[test]
    fn ctas_produces_dag() {
        with_large_stack(|| {
            let sql = "CREATE TABLE new_table AS SELECT id, name FROM source_table";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert_eq!(result.stmt_name, "NEW_TABLE");
            assert!(result.nodes.contains_key(FINAL_SELECT));
            let has_source = result
                .nodes
                .values()
                .any(|n| n.kind == "table" && n.name.as_deref() == Some("source_table"));
            assert!(has_source, "Should find source_table node");
        });
    }

    #[test]
    fn merge_produces_dag() {
        with_large_stack(|| {
            let sql = "MERGE INTO target_tbl USING source_tbl ON target_tbl.id = source_tbl.id \
                        WHEN MATCHED THEN UPDATE SET name = source_tbl.name";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert_eq!(result.stmt_name, "TARGET_TBL");
            assert!(result.nodes.contains_key(FINAL_SELECT));
        });
    }

    #[test]
    fn update_with_from_produces_dag() {
        with_large_stack(|| {
            let sql = "UPDATE target SET col = s.col FROM source s WHERE target.id = s.id";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert_eq!(result.stmt_name, "TARGET");
            assert!(result.nodes.contains_key(FINAL_SELECT));
        });
    }

    #[test]
    fn per_prefix_counters_zero_based() {
        with_large_stack(|| {
            let sql = "SELECT a.id, b.id FROM t1 a \
                        JOIN t2 b ON a.id = b.id \
                        JOIN t3 c ON a.id = c.id";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            let join_ids: Vec<&str> = result
                .nodes
                .values()
                .filter(|n| n.kind.contains("JOIN"))
                .map(|n| n.node_id.as_str())
                .collect();
            assert!(
                join_ids.contains(&"Join_0"),
                "Should have Join_0, got: {join_ids:?}"
            );
            assert!(
                join_ids.contains(&"Join_1"),
                "Should have Join_1, got: {join_ids:?}"
            );
        });
    }

    #[test]
    fn schema_enrichment_populates_columns() {
        with_large_stack(|| {
            let mut schema = HashMap::new();
            let mut cols = HashMap::new();
            cols.insert("ID".to_string(), "NUMBER".to_string());
            cols.insert("NAME".to_string(), "VARCHAR".to_string());
            schema.insert("USERS".to_string(), cols);

            let result = build_dag_with_schema(
                "SELECT id, name FROM users",
                DialectType::Snowflake,
                Some(&schema),
            )
            .unwrap();

            let table_node = result
                .nodes
                .values()
                .find(|n| n.kind == "table" && n.name.as_deref() == Some("users"))
                .expect("Should find users table node");

            assert!(
                !table_node.columns.is_empty(),
                "Table node should have enriched columns"
            );
            let has_id = table_node
                .columns
                .iter()
                .any(|c| c.name == "ID" && c.data_type.as_deref() == Some("NUMBER"));
            assert!(has_id, "Should have ID column with NUMBER type");
        });
    }

    #[test]
    fn slash_comments_stripped() {
        with_large_stack(|| {
            let sql = "// This is a comment\nSELECT id FROM users";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert!(result.nodes.contains_key(FINAL_SELECT));
        });
    }

    #[test]
    fn build_dag_backward_compatible() {
        with_large_stack(|| {
            let sql = "SELECT id, name FROM users WHERE active = true";
            let result = build_dag(sql, DialectType::Snowflake).unwrap();
            assert!(
                result.stmt_name.is_empty(),
                "Plain SELECT should have empty stmt_name"
            );
            assert!(result.nodes.contains_key(FINAL_SELECT));
            let has_table = result.nodes.values().any(|n| n.kind == "table");
            let has_filter = result.nodes.values().any(|n| n.kind == "FILTER");
            assert!(has_table, "Should have table node");
            assert!(has_filter, "Should have filter node");
        });
    }

    #[test]
    fn sqlparser_fallback_three_part_insert() {
        with_large_stack(|| {
            // polyglot-sql fails on 3-part dotted INSERT targets
            let sql = "INSERT INTO catalog.schema.target_table SELECT id FROM source_table";
            let result = build_dag(sql, DialectType::Snowflake);
            // Should succeed via fallback (or polyglot handles it)
            match result {
                Ok(r) => {
                    assert_eq!(r.stmt_name, "TARGET_TABLE");
                }
                Err(e) => {
                    // If both parsers fail, at least the error should be informative
                    assert!(e.contains("fallback"), "Error should mention fallback: {e}");
                }
            }
        });
    }
}
