//! Column Lineage Tracking
//!
//! This module provides functionality to track column lineage through SQL queries,
//! building a graph of how columns flow from source tables to the result set.
//! Supports UNION/INTERSECT/EXCEPT, CTEs, derived tables, subqueries, and star expansion.
//!

use crate::dialects::DialectType;
use crate::expressions::Expression;
use crate::scope::{build_scope, Scope};
use crate::traversal::ExpressionWalk;
use crate::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;

/// A node in the column lineage graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LineageNode {
    /// Name of this lineage step (e.g., "table.column")
    pub name: String,
    /// The expression at this node
    pub expression: Expression,
    /// The source expression (the full query context)
    pub source: Expression,
    /// Downstream nodes that depend on this one
    pub downstream: Vec<LineageNode>,
    /// Optional source name (e.g., for derived tables)
    pub source_name: String,
    /// Optional reference node name (e.g., for CTEs)
    pub reference_node_name: String,
}

impl LineageNode {
    /// Create a new lineage node
    pub fn new(name: impl Into<String>, expression: Expression, source: Expression) -> Self {
        Self {
            name: name.into(),
            expression,
            source,
            downstream: Vec::new(),
            source_name: String::new(),
            reference_node_name: String::new(),
        }
    }

    /// Iterate over all nodes in the lineage graph using DFS
    pub fn walk(&self) -> LineageWalker<'_> {
        LineageWalker { stack: vec![self] }
    }

    /// Get all downstream column names
    pub fn downstream_names(&self) -> Vec<String> {
        self.downstream.iter().map(|n| n.name.clone()).collect()
    }
}

/// Iterator for walking the lineage graph
pub struct LineageWalker<'a> {
    stack: Vec<&'a LineageNode>,
}

impl<'a> Iterator for LineageWalker<'a> {
    type Item = &'a LineageNode;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(node) = self.stack.pop() {
            // Add children in reverse order so they're visited in order
            for child in node.downstream.iter().rev() {
                self.stack.push(child);
            }
            Some(node)
        } else {
            None
        }
    }
}

// ---------------------------------------------------------------------------
// ColumnRef: name or positional index for column lookup
// ---------------------------------------------------------------------------

/// Column reference for lineage tracing — by name or positional index.
enum ColumnRef<'a> {
    Name(&'a str),
    Index(usize),
}

// ---------------------------------------------------------------------------
// Lineage options
// ---------------------------------------------------------------------------

/// Configuration for column lineage tracing.
#[derive(Debug, Clone, Default)]
pub struct LineageOptions {
    /// When true, only trace value-producing column references.
    /// Skips CASE WHEN conditions, IFF condition arg, and window
    /// PARTITION BY / ORDER BY columns.  Matches Snowflake ACCESS_USAGE
    /// "direct" lineage semantics where condition columns are NOT
    /// considered lineage sources.
    pub values_only: bool,
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Build the lineage graph for a single output column in a parsed SQL AST.
///
/// Traces the column backward through CTEs, derived tables, set operations,
/// and expressions to find the source table columns that contribute to it.
///
/// # Arguments
/// * `column` - The output column name to trace lineage for
/// * `sql` - The parsed SQL expression (SELECT, UNION, etc.)
/// * `dialect` - Optional dialect for identifier resolution
/// * `trim_selects` - If true, trim the source SELECT to only the target column
/// * `options` - Controls lineage behavior (e.g., `values_only` to skip conditions)
///
/// # Returns
/// The root [`LineageNode`] with downstream nodes forming a lineage graph.
///
/// # Example
/// ```ignore
/// use polyglot_sql::lineage::{lineage, LineageOptions};
/// use polyglot_sql::parse_one;
/// use polyglot_sql::DialectType;
///
/// let sql = "SELECT a, b + 1 AS c FROM t";
/// let expr = parse_one(sql, DialectType::Generic).unwrap();
/// let opts = LineageOptions { values_only: true };
/// let node = lineage("c", &expr, None, false, Some(&opts)).unwrap();
/// ```
pub fn lineage(
    column: &str,
    sql: &Expression,
    dialect: Option<DialectType>,
    trim_selects: bool,
    options: Option<&LineageOptions>,
) -> Result<LineageNode> {
    let defaults = LineageOptions::default();
    let opts = options.unwrap_or(&defaults);
    let scope = build_scope(sql);
    to_node(
        ColumnRef::Name(column),
        &scope,
        dialect,
        "",
        "",
        "",
        trim_selects,
        opts.values_only,
    )
}

/// Get all source tables from a lineage graph
pub fn get_source_tables(node: &LineageNode) -> HashSet<String> {
    let mut tables = HashSet::new();
    collect_source_tables(node, &mut tables);
    tables
}

/// Recursively collect source table names from lineage graph
pub fn collect_source_tables(node: &LineageNode, tables: &mut HashSet<String>) {
    if let Expression::Table(table) = &node.source {
        tables.insert(table.name.name.clone());
    }
    for child in &node.downstream {
        collect_source_tables(child, tables);
    }
}

// ---------------------------------------------------------------------------
// Core recursive lineage builder
// ---------------------------------------------------------------------------

/// Maximum recursion depth for lineage tracing.  Deeply nested derived tables
/// and recursive CTEs can cause unbounded recursion → stack overflow.  When the
/// limit is hit we return a placeholder terminal instead of crashing.
const MAX_LINEAGE_DEPTH: usize = 50;

/// Recursively build a lineage node for a column in a scope.
fn to_node(
    column: ColumnRef<'_>,
    scope: &Scope,
    dialect: Option<DialectType>,
    scope_name: &str,
    source_name: &str,
    reference_node_name: &str,
    trim_selects: bool,
    values_only: bool,
) -> Result<LineageNode> {
    to_node_inner(
        column,
        scope,
        dialect,
        scope_name,
        source_name,
        reference_node_name,
        trim_selects,
        &[],
        0,
        values_only,
    )
}

fn to_node_inner(
    column: ColumnRef<'_>,
    scope: &Scope,
    dialect: Option<DialectType>,
    scope_name: &str,
    source_name: &str,
    reference_node_name: &str,
    trim_selects: bool,
    ancestor_cte_scopes: &[&Scope],
    depth: usize,
    values_only: bool,
) -> Result<LineageNode> {
    // Guard against unbounded recursion (recursive CTEs, deeply nested subqueries).
    if depth > MAX_LINEAGE_DEPTH {
        return Err(crate::error::Error::Parse(
            "Column lineage recursion depth exceeded".to_string(),
        ));
    }
    let scope_expr = &scope.expression;

    // Build combined CTE scopes: current scope's cte_scopes + ancestors.
    // All references — zero cloning.
    let mut all_cte_scopes: Vec<&Scope> = scope.cte_scopes.iter().collect();
    for &s in ancestor_cte_scopes {
        all_cte_scopes.push(s);
    }

    // 0. Unwrap CTE scope — CTE scope expressions are Expression::Cte(...)
    //    but we need the inner query (SELECT/UNION) for column lookup.
    let effective_expr = match scope_expr {
        Expression::Cte(cte) => &cte.this,
        other => other,
    };

    // 1. Set operations (UNION / INTERSECT / EXCEPT)
    if matches!(
        effective_expr,
        Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
    ) {
        // For CTE wrapping a set op, create a temporary scope with the inner expression
        if matches!(scope_expr, Expression::Cte(_)) {
            let mut inner_scope = Scope::new(effective_expr.clone());
            inner_scope.union_scopes = scope.union_scopes.clone();
            inner_scope.sources = scope.sources.clone();
            inner_scope.cte_sources = scope.cte_sources.clone();
            inner_scope.cte_scopes = scope.cte_scopes.clone();
            inner_scope.derived_table_scopes = scope.derived_table_scopes.clone();
            inner_scope.subquery_scopes = scope.subquery_scopes.clone();
            return handle_set_operation(
                &column,
                &inner_scope,
                dialect,
                scope_name,
                source_name,
                reference_node_name,
                trim_selects,
                ancestor_cte_scopes,
                depth,
                values_only,
            );
        }
        return handle_set_operation(
            &column,
            scope,
            dialect,
            scope_name,
            source_name,
            reference_node_name,
            trim_selects,
            ancestor_cte_scopes,
            depth,
            values_only,
        );
    }

    // 2. Find the select expression for this column
    let select_expr = find_select_expr(effective_expr, &column)?;
    let column_name = resolve_column_name(&column, &select_expr);

    // 3. Trim source if requested
    let node_source = if trim_selects {
        trim_source(effective_expr, &select_expr)
    } else {
        effective_expr.clone()
    };

    // 4. Create the lineage node
    let mut node = LineageNode::new(&column_name, select_expr.clone(), node_source);
    node.source_name = source_name.to_string();
    node.reference_node_name = reference_node_name.to_string();

    // 5. Star handling — expand through CTE / derived-table scopes, or
    //    fall back to a star terminal for raw tables.
    if matches!(&select_expr, Expression::Star(_)) {
        for (name, source_info) in &scope.sources {
            // CTE source — resolve individual columns from the CTE definition.
            if scope.cte_sources.contains_key(name) {
                if let Some(child_scope) = find_child_scope_in(&all_cte_scopes, scope, name) {
                    let cte_cols = extract_output_columns(&child_scope.expression);
                    if !cte_cols.is_empty() && !cte_cols.contains(&"*".to_string()) {
                        for col in &cte_cols {
                            if let Ok(child) = to_node_inner(
                                ColumnRef::Name(col),
                                child_scope,
                                dialect,
                                scope_name,
                                name,
                                scope_name,
                                trim_selects,
                                &all_cte_scopes,
                                depth + 1,
                                values_only,
                            ) {
                                node.downstream.push(child);
                            }
                        }
                        continue;
                    }
                }
            }

            // Derived table source — resolve individual columns.
            if source_info.is_scope {
                if let Some(child_scope) = find_child_scope(scope, name) {
                    let sub_cols = extract_output_columns(&child_scope.expression);
                    if !sub_cols.is_empty() && !sub_cols.contains(&"*".to_string()) {
                        for col in &sub_cols {
                            if let Ok(child) = to_node_inner(
                                ColumnRef::Name(col),
                                child_scope,
                                dialect,
                                scope_name,
                                name,
                                scope_name,
                                trim_selects,
                                &all_cte_scopes,
                                depth + 1,
                                values_only,
                            ) {
                                node.downstream.push(child);
                            }
                        }
                        continue;
                    }
                }
            }

            // Raw table or unresolvable — keep as star terminal.
            let child = LineageNode::new(
                format!("{}.*", name),
                Expression::Star(crate::expressions::Star {
                    table: None,
                    except: None,
                    replace: None,
                    rename: None,
                    trailing_comments: vec![],
                }),
                source_info.expression.clone(),
            );
            node.downstream.push(child);
        }
        return Ok(node);
    }

    // 6. Subqueries in select — trace through scalar subqueries
    let subqueries: Vec<&Expression> =
        select_expr.find_all(|e| matches!(e, Expression::Subquery(sq) if sq.alias.is_none()));
    for sq_expr in subqueries {
        if let Expression::Subquery(sq) = sq_expr {
            for sq_scope in &scope.subquery_scopes {
                if sq_scope.expression == sq.this {
                    if let Ok(child) = to_node_inner(
                        ColumnRef::Index(0),
                        sq_scope,
                        dialect,
                        &column_name,
                        "",
                        "",
                        trim_selects,
                        &all_cte_scopes,
                        depth + 1,
                        values_only,
                    ) {
                        node.downstream.push(child);
                    }
                    break;
                }
            }
        }
    }

    // 7. Column references — trace each column to its source
    let col_refs = find_column_refs_in_expr_opts(&select_expr, values_only);
    for col_ref in col_refs {
        let col_name = &col_ref.column;
        if let Some(ref table_id) = col_ref.table {
            let tbl = &table_id.name;
            resolve_qualified_column(
                &mut node,
                scope,
                dialect,
                tbl,
                col_name,
                &column_name,
                trim_selects,
                &all_cte_scopes,
                depth,
                values_only,
            );
        } else {
            resolve_unqualified_column(
                &mut node,
                scope,
                dialect,
                col_name,
                &column_name,
                trim_selects,
                &all_cte_scopes,
                depth,
                values_only,
            );
        }
    }

    Ok(node)
}

// ---------------------------------------------------------------------------
// Set operation handling
// ---------------------------------------------------------------------------

fn handle_set_operation(
    column: &ColumnRef<'_>,
    scope: &Scope,
    dialect: Option<DialectType>,
    scope_name: &str,
    source_name: &str,
    reference_node_name: &str,
    trim_selects: bool,
    ancestor_cte_scopes: &[&Scope],
    depth: usize,
    values_only: bool,
) -> Result<LineageNode> {
    let scope_expr = &scope.expression;

    // Determine column index
    let col_index = match column {
        ColumnRef::Name(name) => column_to_index(scope_expr, name)?,
        ColumnRef::Index(i) => *i,
    };

    let col_name = match column {
        ColumnRef::Name(name) => name.to_string(),
        ColumnRef::Index(_) => format!("_{col_index}"),
    };

    let mut node = LineageNode::new(&col_name, scope_expr.clone(), scope_expr.clone());
    node.source_name = source_name.to_string();
    node.reference_node_name = reference_node_name.to_string();

    // Recurse into each union branch
    for branch_scope in &scope.union_scopes {
        if let Ok(child) = to_node_inner(
            ColumnRef::Index(col_index),
            branch_scope,
            dialect,
            scope_name,
            "",
            "",
            trim_selects,
            ancestor_cte_scopes,
            depth + 1,
            values_only,
        ) {
            node.downstream.push(child);
        }
    }

    Ok(node)
}

// ---------------------------------------------------------------------------
// Column resolution helpers
// ---------------------------------------------------------------------------

fn resolve_qualified_column(
    node: &mut LineageNode,
    scope: &Scope,
    dialect: Option<DialectType>,
    table: &str,
    col_name: &str,
    parent_name: &str,
    trim_selects: bool,
    all_cte_scopes: &[&Scope],
    depth: usize,
    values_only: bool,
) {
    // Check if table is a CTE reference (cte_sources tracks CTE names)
    if scope.cte_sources.contains_key(table) {
        if let Some(child_scope) = find_child_scope_in(all_cte_scopes, scope, table) {
            if let Ok(child) = to_node_inner(
                ColumnRef::Name(col_name),
                child_scope,
                dialect,
                parent_name,
                table,
                parent_name,
                trim_selects,
                all_cte_scopes,
                depth + 1,
                values_only,
            ) {
                node.downstream.push(child);
                return;
            }
        }
    }

    // Check if table is a derived table (is_scope = true in sources)
    if let Some(source_info) = scope.sources.get(table) {
        if source_info.is_scope {
            if let Some(child_scope) = find_child_scope(scope, table) {
                if let Ok(child) = to_node_inner(
                    ColumnRef::Name(col_name),
                    child_scope,
                    dialect,
                    parent_name,
                    table,
                    parent_name,
                    trim_selects,
                    all_cte_scopes,
                    depth + 1,
                    values_only,
                ) {
                    node.downstream.push(child);
                    return;
                }
            }
        }
    }

    // Base table or unresolved — terminal node (preserve full table path)
    let source_expr = scope.sources.get(table).map(|si| &si.expression);
    node.downstream
        .push(make_table_column_node(table, col_name, source_expr));
}

fn resolve_unqualified_column(
    node: &mut LineageNode,
    scope: &Scope,
    dialect: Option<DialectType>,
    col_name: &str,
    parent_name: &str,
    trim_selects: bool,
    all_cte_scopes: &[&Scope],
    depth: usize,
    values_only: bool,
) {
    // Try to find which source this column belongs to.
    // Filter to only FROM-clause sources: add_cte_source adds all CTEs to sources
    // with Expression::Cte, but FROM-clause Table references overwrite with Expression::Table.
    // So CTE-only entries (not referenced in FROM) have Expression::Cte — exclude those.
    let from_source_names: Vec<&String> = scope
        .sources
        .iter()
        .filter(|(_, info)| !matches!(info.expression, Expression::Cte(_)))
        .map(|(name, _)| name)
        .collect();

    if from_source_names.len() == 1 {
        let tbl = from_source_names[0];
        resolve_qualified_column(
            node,
            scope,
            dialect,
            tbl,
            col_name,
            parent_name,
            trim_selects,
            all_cte_scopes,
            depth,
            values_only,
        );
        return;
    }

    // Multiple sources — can't resolve without schema info, add unqualified node
    let child = LineageNode::new(
        col_name.to_string(),
        Expression::Column(crate::expressions::Column {
            name: crate::expressions::Identifier::new(col_name.to_string()),
            table: None,
            join_mark: false,
            trailing_comments: vec![],
        }),
        node.source.clone(),
    );
    node.downstream.push(child);
}

// ---------------------------------------------------------------------------
// Helper functions
// ---------------------------------------------------------------------------

/// Get the alias or name of an expression
fn get_alias_or_name(expr: &Expression) -> Option<String> {
    match expr {
        Expression::Alias(alias) => Some(alias.alias.name.clone()),
        Expression::Column(col) => Some(col.name.name.clone()),
        Expression::Identifier(id) => Some(id.name.clone()),
        Expression::Star(_) => Some("*".to_string()),
        _ => None,
    }
}

/// Resolve the display name for a column reference.
fn resolve_column_name(column: &ColumnRef<'_>, select_expr: &Expression) -> String {
    match column {
        ColumnRef::Name(n) => n.to_string(),
        ColumnRef::Index(_) => get_alias_or_name(select_expr).unwrap_or_else(|| "?".to_string()),
    }
}

/// Find the select expression matching a column reference.
fn find_select_expr(scope_expr: &Expression, column: &ColumnRef<'_>) -> Result<Expression> {
    if let Expression::Select(ref select) = scope_expr {
        match column {
            ColumnRef::Name(name) => {
                for expr in &select.expressions {
                    if get_alias_or_name(expr).as_deref() == Some(name) {
                        return Ok(expr.clone());
                    }
                }
                Err(crate::error::Error::Parse(format!(
                    "Cannot find column '{}' in query",
                    name
                )))
            }
            ColumnRef::Index(idx) => select.expressions.get(*idx).cloned().ok_or_else(|| {
                crate::error::Error::Parse(format!("Column index {} out of range", idx))
            }),
        }
    } else {
        Err(crate::error::Error::Parse(
            "Expected SELECT expression for column lookup".to_string(),
        ))
    }
}

/// Find the positional index of a column name in a set operation's first SELECT branch.
fn column_to_index(set_op_expr: &Expression, name: &str) -> Result<usize> {
    let mut expr = set_op_expr;
    loop {
        match expr {
            Expression::Union(u) => expr = &u.left,
            Expression::Intersect(i) => expr = &i.left,
            Expression::Except(e) => expr = &e.left,
            Expression::Select(select) => {
                for (i, e) in select.expressions.iter().enumerate() {
                    if get_alias_or_name(e).as_deref() == Some(name) {
                        return Ok(i);
                    }
                }
                return Err(crate::error::Error::Parse(format!(
                    "Cannot find column '{}' in set operation",
                    name
                )));
            }
            _ => {
                return Err(crate::error::Error::Parse(
                    "Expected SELECT or set operation".to_string(),
                ))
            }
        }
    }
}

/// If trim_selects is enabled, return a copy of the SELECT with only the target column.
fn trim_source(select_expr: &Expression, target_expr: &Expression) -> Expression {
    if let Expression::Select(select) = select_expr {
        let mut trimmed = select.as_ref().clone();
        trimmed.expressions = vec![target_expr.clone()];
        Expression::Select(Box::new(trimmed))
    } else {
        select_expr.clone()
    }
}

/// Find the child scope (CTE or derived table) for a given source name.
fn find_child_scope<'a>(scope: &'a Scope, source_name: &str) -> Option<&'a Scope> {
    // Check CTE scopes
    if scope.cte_sources.contains_key(source_name) {
        for cte_scope in &scope.cte_scopes {
            if let Expression::Cte(cte) = &cte_scope.expression {
                if cte.alias.name == source_name {
                    return Some(cte_scope);
                }
            }
        }
    }

    // Check derived table scopes
    if let Some(source_info) = scope.sources.get(source_name) {
        if source_info.is_scope && !scope.cte_sources.contains_key(source_name) {
            if let Expression::Subquery(sq) = &source_info.expression {
                for dt_scope in &scope.derived_table_scopes {
                    if dt_scope.expression == sq.this {
                        return Some(dt_scope);
                    }
                }
            }
        }
    }

    None
}

/// Find a CTE scope by name, searching through a combined list of CTE scopes.
/// This handles nested CTEs where the current scope doesn't have the CTE scope
/// as a direct child but knows about it via cte_sources.
fn find_child_scope_in<'a>(
    all_cte_scopes: &[&'a Scope],
    scope: &'a Scope,
    source_name: &str,
) -> Option<&'a Scope> {
    // First try the scope's own cte_scopes
    for cte_scope in &scope.cte_scopes {
        if let Expression::Cte(cte) = &cte_scope.expression {
            if cte.alias.name == source_name {
                return Some(cte_scope);
            }
        }
    }

    // Then search through all ancestor CTE scopes
    for cte_scope in all_cte_scopes {
        if let Expression::Cte(cte) = &cte_scope.expression {
            if cte.alias.name == source_name {
                return Some(cte_scope);
            }
        }
    }

    // Fall back to derived table scopes
    if let Some(source_info) = scope.sources.get(source_name) {
        if source_info.is_scope {
            if let Expression::Subquery(sq) = &source_info.expression {
                for dt_scope in &scope.derived_table_scopes {
                    if dt_scope.expression == sq.this {
                        return Some(dt_scope);
                    }
                }
            }
        }
    }

    None
}

/// Create a terminal lineage node for a table.column reference.
///
/// When `source_table_expr` is provided, it's used as the source — preserving
/// the full catalog.schema.table path. Otherwise falls back to a simple TableRef.
fn make_table_column_node(
    table: &str,
    column: &str,
    source_table_expr: Option<&Expression>,
) -> LineageNode {
    let source = source_table_expr
        .cloned()
        .unwrap_or_else(|| Expression::Table(crate::expressions::TableRef::new(table)));
    LineageNode::new(
        format!("{}.{}", table, column),
        Expression::Column(crate::expressions::Column {
            name: crate::expressions::Identifier::new(column.to_string()),
            table: Some(crate::expressions::Identifier::new(table.to_string())),
            join_mark: false,
            trailing_comments: vec![],
        }),
        source,
    )
}

/// Simple column reference extracted from an expression
#[derive(Debug, Clone)]
struct SimpleColumnRef {
    table: Option<crate::expressions::Identifier>,
    column: String,
}

/// Find column references in an expression with values-only filtering.
fn find_column_refs_in_expr_opts(expr: &Expression, values_only: bool) -> Vec<SimpleColumnRef> {
    let mut refs = Vec::new();
    collect_column_refs(expr, &mut refs, values_only);
    refs
}

fn collect_column_refs(expr: &Expression, refs: &mut Vec<SimpleColumnRef>, values_only: bool) {
    match expr {
        Expression::Column(col) => {
            refs.push(SimpleColumnRef {
                table: col.table.clone(),
                column: col.name.name.clone(),
            });
        }
        Expression::Alias(alias) => {
            collect_column_refs(&alias.this, refs, values_only);
        }
        Expression::And(op)
        | Expression::Or(op)
        | Expression::Eq(op)
        | Expression::Neq(op)
        | Expression::Lt(op)
        | Expression::Lte(op)
        | Expression::Gt(op)
        | Expression::Gte(op)
        | Expression::Add(op)
        | Expression::Sub(op)
        | Expression::Mul(op)
        | Expression::Div(op)
        | Expression::Mod(op)
        | Expression::BitwiseAnd(op)
        | Expression::BitwiseOr(op)
        | Expression::BitwiseXor(op)
        | Expression::Concat(op) => {
            collect_column_refs(&op.left, refs, values_only);
            collect_column_refs(&op.right, refs, values_only);
        }
        Expression::Not(u) | Expression::Neg(u) | Expression::BitwiseNot(u) => {
            collect_column_refs(&u.this, refs, values_only);
        }
        Expression::Function(func) => {
            let func_name_upper = func.name.to_uppercase();

            // IFF(cond, then, else) / DECODE(expr, val, result, ...):
            // In values_only mode, skip the condition (first) argument.
            let skip_first =
                values_only && (func_name_upper == "IFF" || func_name_upper == "DECODE");

            // DATE_PART, EXTRACT, DATEDIFF, etc.: first arg is a date-part keyword,
            // not a column reference (e.g., EPOCH_MILLISECOND, YEAR, MONTH).
            static DATE_PART_FUNCS: &[&str] = &[
                "DATE_PART",
                "DATEPART",
                "EXTRACT",
                "DATE_TRUNC",
                "DATETRUNC",
                "TRUNC",
                "DATEDIFF",
                "DATEADD",
                "DATE_ADD",
                "DATE_SUB",
                "TIMESTAMPADD",
                "TIMESTAMPDIFF",
                "TIME_SLICE",
            ];
            let skip_first_keyword = DATE_PART_FUNCS.contains(&func_name_upper.as_str());

            for (idx, arg) in func.args.iter().enumerate() {
                if idx == 0 && (skip_first || skip_first_keyword) {
                    continue;
                }
                collect_column_refs(arg, refs, values_only);
            }
        }
        Expression::AggregateFunction(func) => {
            for arg in &func.args {
                collect_column_refs(arg, refs, values_only);
            }
        }
        Expression::WindowFunction(wf) => {
            // In values_only mode, only trace the function expression itself.
            // Snowflake ACCESS_USAGE excludes PARTITION BY and ORDER BY from
            // column lineage — they control grouping/ordering, not data flow.
            collect_column_refs(&wf.this, refs, values_only);
        }
        Expression::Case(case) => {
            if values_only {
                // Values-only: skip operand and conditions, keep results and else.
                for (_cond, result) in &case.whens {
                    collect_column_refs(result, refs, values_only);
                }
                if let Some(ref else_expr) = case.else_ {
                    collect_column_refs(else_expr, refs, values_only);
                }
            } else {
                if let Some(operand) = &case.operand {
                    collect_column_refs(operand, refs, values_only);
                }
                for (cond, result) in &case.whens {
                    collect_column_refs(cond, refs, values_only);
                    collect_column_refs(result, refs, values_only);
                }
                if let Some(ref else_expr) = case.else_ {
                    collect_column_refs(else_expr, refs, values_only);
                }
            }
        }
        Expression::Cast(cast) => {
            collect_column_refs(&cast.this, refs, values_only);
        }
        Expression::Paren(p) => {
            collect_column_refs(&p.this, refs, values_only);
        }
        Expression::Coalesce(c) => {
            for e in &c.expressions {
                collect_column_refs(e, refs, values_only);
            }
        }
        // UnaryFunc variants — these are dedicated expression types for common
        // SQL functions (UPPER, LOWER, etc.) that wrap a single child expression.
        // They must be listed explicitly because the traversal framework's
        // iter_children/children() doesn't handle them.
        Expression::Upper(f)
        | Expression::Lower(f)
        | Expression::Length(f)
        | Expression::LTrim(f)
        | Expression::RTrim(f)
        | Expression::Reverse(f)
        | Expression::Abs(f)
        | Expression::Sqrt(f)
        | Expression::Cbrt(f)
        | Expression::Ln(f)
        | Expression::Exp(f)
        | Expression::Sign(f)
        | Expression::Date(f)
        | Expression::Time(f)
        | Expression::Year(f)
        | Expression::Month(f)
        | Expression::Day(f)
        | Expression::Hour(f)
        | Expression::Minute(f)
        | Expression::Second(f)
        | Expression::DayOfWeek(f)
        | Expression::DayOfMonth(f)
        | Expression::DayOfYear(f)
        | Expression::WeekOfYear(f)
        | Expression::Quarter(f)
        | Expression::Epoch(f)
        | Expression::Initcap(f)
        | Expression::Ascii(f)
        | Expression::Chr(f)
        | Expression::Soundex(f)
        | Expression::ByteLength(f)
        | Expression::Hex(f)
        | Expression::Radians(f)
        | Expression::Degrees(f)
        | Expression::Sin(f)
        | Expression::Cos(f)
        | Expression::Tan(f)
        | Expression::Asin(f)
        | Expression::Acos(f)
        | Expression::Atan(f)
        | Expression::IsNan(f)
        | Expression::IsInf(f)
        | Expression::Typeof(f)
        | Expression::ArrayLength(f)
        | Expression::ArraySize(f)
        | Expression::Cardinality(f)
        | Expression::ArrayReverse(f)
        | Expression::ArrayDistinct(f)
        | Expression::Explode(f)
        | Expression::ExplodeOuter(f)
        | Expression::ArrayFlatten(f)
        | Expression::ArrayCompact(f)
        | Expression::ToArray(f)
        | Expression::MapFromEntries(f)
        | Expression::MapKeys(f)
        | Expression::MapValues(f)
        | Expression::JsonArrayLength(f)
        | Expression::JsonKeys(f)
        | Expression::JsonType(f)
        | Expression::ParseJson(f)
        | Expression::ToJson(f)
        | Expression::BitwiseCount(f)
        | Expression::DateFromUnixDate(f)
        | Expression::UnixDate(f)
        | Expression::UnixSeconds(f)
        | Expression::UnixMillis(f)
        | Expression::UnixMicros(f)
        | Expression::TimeStrToDate(f)
        | Expression::DateToDi(f)
        | Expression::DiToDate(f)
        | Expression::TsOrDiToDi(f)
        | Expression::TsOrDsToDatetime(f)
        | Expression::TsOrDsToTimestamp(f)
        | Expression::YearOfWeek(f)
        | Expression::YearOfWeekIso(f)
        | Expression::EpochMs(f)
        | Expression::TimeStrToUnix(f)
        | Expression::TimeToUnix(f)
        | Expression::SHA(f)
        | Expression::SHA1Digest(f)
        | Expression::DateStrToDate(f)
        | Expression::DateToDateStr(f)
        | Expression::Unicode(f)
        | Expression::LowerHex(f)
        | Expression::Int64(f)
        | Expression::JSONBool(f)
        | Expression::MD5NumberLower64(f)
        | Expression::MD5NumberUpper64(f) => {
            collect_column_refs(&f.this, refs, values_only);
        }
        // Trim has a different struct but also wraps an expression
        Expression::Trim(t) => {
            collect_column_refs(&t.this, refs, values_only);
        }
        // Don't recurse into subqueries — those are handled separately
        Expression::Subquery(_) | Expression::Exists(_) => {}
        _ => {}
    }
}

// ===========================================================================
// Column Lineage — full-query edge extraction
// ===========================================================================

/// How a column's value was transformed between source and target.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum LensType {
    /// Column passes through unmodified.
    Unchanged,
    /// Column was renamed (simple alias of a single column).
    Alias,
    /// Column was computed from an expression.
    Transformation,
}

impl LensType {
    pub fn as_str(&self) -> &'static str {
        match self {
            LensType::Unchanged => "Unchanged",
            LensType::Alias => "Alias",
            LensType::Transformation => "Transformation",
        }
    }
}

/// What kind of terminal the lineage trace ended at.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum TerminalKind {
    /// Resolved to a concrete table.column.
    Table,
    /// Resolved to a SELECT * (table inferred from FROM).
    Star,
    /// Parser returned a placeholder (ambiguous star / aggregate).
    Placeholder,
    /// Could not resolve the terminal.
    Unknown,
}

impl TerminalKind {
    pub fn as_str(&self) -> &'static str {
        match self {
            TerminalKind::Table => "table",
            TerminalKind::Star => "star",
            TerminalKind::Placeholder => "placeholder",
            TerminalKind::Unknown => "unknown",
        }
    }
}

/// One step in the transformation lens (the chain of expressions from
/// target column back to source column).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LensStep {
    /// SQL text of the expression at this step.
    pub expression_sql: String,
    /// Classification: "passthrough", "alias", "transform", or the terminal
    /// kind ("table", "star", "placeholder").
    pub code_type: String,
}

/// A single source→target column lineage edge.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ColumnLineageEdge {
    /// Dot-joined source table path: "catalog.schema.table" (unquoted).
    pub source_table: String,
    /// Source column name.
    pub source_column: String,
    /// Target column name (in the outermost SELECT).
    pub target_column: String,
    /// Edge type — always "direct" for intra-query lineage.
    pub edge_type: String,
    /// How the value was transformed.
    pub lens_type: LensType,
    /// The full transformation chain, ordered source→target.
    pub lens_code: Vec<LensStep>,
    /// What kind of terminal the trace ended at.
    pub terminal_kind: TerminalKind,
}

/// An error encountered while tracing a single column.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ColumnLineageError {
    pub column: String,
    pub message: String,
}

/// Result of running column lineage over an entire query.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ColumnLineageResult {
    pub edges: Vec<ColumnLineageEdge>,
    pub errors: Vec<ColumnLineageError>,
}

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Extract column lineage edges for every output column in a SQL query.
///
/// This is the main entry point for column-level lineage analysis. It:
/// 1. Parses the SQL into an AST using the given dialect
/// 2. Identifies all output columns from the outermost SELECT
/// 3. Traces each column backward to its source table columns
/// 4. Returns edges with transformation metadata (LensType, lens_code)
///
/// # Arguments
/// * `sql` - The SQL query string to analyze
/// * `dialect` - The SQL dialect for parsing (e.g., Snowflake, BigQuery)
/// * `options` - Controls lineage behavior:
///   - `values_only: true` → skip CASE conditions, IFF condition args,
///     date-part keywords, window PARTITION BY/ORDER BY
///     (matches Snowflake ACCESS_USAGE "direct" lineage semantics)
///   - `values_only: false` → trace ALL column references including conditions
///
/// # Returns
/// [`ColumnLineageResult`] containing:
/// - `edges`: One [`ColumnLineageEdge`] per source→target column mapping, with:
///   - `source_table` / `source_column`: where the data comes from
///   - `target_column`: output column name
///   - `lens_type`: Unchanged / Alias / Transformation
///   - `lens_code`: full transformation chain (source→target)
///   - `terminal_kind`: Table / Star / Placeholder / Unknown
/// - `errors`: columns that couldn't be traced (parse failures, ambiguity)
pub fn column_lineage(
    sql: &str,
    dialect: DialectType,
    options: Option<&LineageOptions>,
) -> ColumnLineageResult {
    let defaults = LineageOptions::default();
    let opts = options.unwrap_or(&defaults);
    let d = crate::dialects::Dialect::get(dialect);
    let exprs = match d.parse(sql) {
        Ok(e) => e,
        Err(e) => {
            return ColumnLineageResult {
                edges: vec![],
                errors: vec![ColumnLineageError {
                    column: "*".into(),
                    message: format!("Parse error: {e}"),
                }],
            };
        }
    };
    let first = match exprs.into_iter().next() {
        Some(e) => e,
        None => {
            return ColumnLineageResult {
                edges: vec![],
                errors: vec![ColumnLineageError {
                    column: "*".into(),
                    message: "No SQL statements found".into(),
                }],
            };
        }
    };

    let dialect_opt = if dialect == DialectType::Generic {
        None
    } else {
        Some(dialect)
    };

    // Single-pass: parse once, build scope once, trace all columns in one
    // traversal.  No per-column `lineage()` loop.
    let scope = build_scope(&first);
    let mut result = ColumnLineageResult {
        edges: vec![],
        errors: vec![],
    };

    trace_all_columns(&scope, dialect_opt, opts.values_only, &mut result);

    // Reverse lens_code on all edges so they read source → target.
    for edge in &mut result.edges {
        edge.lens_code.reverse();
    }

    result
}

/// Single-pass column lineage: trace ALL output columns from a scope at once.
///
/// Instead of looping over column names and calling `to_node()` per column
/// (which repeats CTE scope resolution and SELECT-list scans), this function
/// directly iterates the SELECT expression list by index, processing each
/// output column in a single walk.
fn trace_all_columns(
    scope: &Scope,
    dialect: Option<DialectType>,
    values_only: bool,
    result: &mut ColumnLineageResult,
) {
    let scope_expr = &scope.expression;

    // Unwrap CTE wrapper if present.
    let effective_expr = match scope_expr {
        Expression::Cte(cte) => &cte.this,
        other => other,
    };

    // Set operations: delegate to per-branch tracing by index.
    if matches!(
        effective_expr,
        Expression::Union(_) | Expression::Intersect(_) | Expression::Except(_)
    ) {
        let col_names = extract_output_columns(effective_expr);
        // For set ops, CTE wrapper needs scope adjustment.
        let effective_scope = if matches!(scope_expr, Expression::Cte(_)) {
            let mut inner = Scope::new(effective_expr.clone());
            inner.union_scopes = scope.union_scopes.clone();
            inner.sources = scope.sources.clone();
            inner.cte_sources = scope.cte_sources.clone();
            inner.cte_scopes = scope.cte_scopes.clone();
            inner.derived_table_scopes = scope.derived_table_scopes.clone();
            inner.subquery_scopes = scope.subquery_scopes.clone();
            Some(inner)
        } else {
            None
        };
        let s = effective_scope.as_ref().unwrap_or(scope);
        for (i, col_name) in col_names.iter().enumerate() {
            match to_node_inner(
                ColumnRef::Index(i),
                s,
                dialect,
                "",
                "",
                "",
                false,
                &[],
                0,
                values_only,
            ) {
                Ok(root) => {
                    walk_lineage_tree(
                        &root,
                        col_name,
                        LensType::Unchanged,
                        vec![],
                        &mut result.edges,
                    );
                }
                Err(e) => {
                    result.errors.push(ColumnLineageError {
                        column: col_name.clone(),
                        message: format!("{e}"),
                    });
                }
            }
        }
        return;
    }

    // Regular SELECT: iterate expressions directly by index.
    if let Expression::Select(ref select) = effective_expr {
        let all_cte_scopes: Vec<&Scope> = scope.cte_scopes.iter().collect();
        for (i, expr) in select.expressions.iter().enumerate() {
            let col_name = get_alias_or_name(expr).unwrap_or_else(|| format!("_{i}"));
            // Directly process the expression — no find_select_expr name-scan.
            match trace_single_column(
                i,
                &col_name,
                expr,
                effective_expr,
                scope,
                dialect,
                &all_cte_scopes,
                values_only,
            ) {
                Ok(root) => {
                    walk_lineage_tree(
                        &root,
                        &col_name,
                        LensType::Unchanged,
                        vec![],
                        &mut result.edges,
                    );
                }
                Err(e) => {
                    result.errors.push(ColumnLineageError {
                        column: col_name,
                        message: format!("{e}"),
                    });
                }
            }
        }
    }
}

/// Trace lineage for a single SELECT expression that we already have in hand.
///
/// This is the core of the single-pass approach: we already know the expression
/// and its output name/index, so we skip the column-name lookup that
/// `to_node_inner` performs via `find_select_expr`.
fn trace_single_column(
    _index: usize,
    col_name: &str,
    select_expr: &Expression,
    scope_expr: &Expression,
    scope: &Scope,
    dialect: Option<DialectType>,
    all_cte_scopes: &[&Scope],
    values_only: bool,
) -> Result<LineageNode> {
    let mut node = LineageNode::new(col_name, select_expr.clone(), scope_expr.clone());

    // Star handling — expand through CTE / derived-table scopes.
    if matches!(select_expr, Expression::Star(_)) {
        expand_star_into_node(&mut node, scope, dialect, all_cte_scopes, values_only);
        return Ok(node);
    }

    // Subqueries in SELECT — trace through scalar subqueries.
    let subqueries: Vec<&Expression> =
        select_expr.find_all(|e| matches!(e, Expression::Subquery(sq) if sq.alias.is_none()));
    for sq_expr in subqueries {
        if let Expression::Subquery(sq) = sq_expr {
            for sq_scope in &scope.subquery_scopes {
                if sq_scope.expression == sq.this {
                    if let Ok(child) = to_node_inner(
                        ColumnRef::Index(0),
                        sq_scope,
                        dialect,
                        col_name,
                        "",
                        "",
                        false,
                        all_cte_scopes,
                        1,
                        values_only,
                    ) {
                        node.downstream.push(child);
                    }
                    break;
                }
            }
        }
    }

    // Column references — trace each column to its source.
    let col_refs = find_column_refs_in_expr_opts(select_expr, values_only);
    for col_ref in col_refs {
        let ref_col = &col_ref.column;
        if let Some(ref table_id) = col_ref.table {
            let tbl = &table_id.name;
            resolve_qualified_column(
                &mut node,
                scope,
                dialect,
                tbl,
                ref_col,
                col_name,
                false,
                all_cte_scopes,
                0,
                values_only,
            );
        } else {
            resolve_unqualified_column(
                &mut node,
                scope,
                dialect,
                ref_col,
                col_name,
                false,
                all_cte_scopes,
                0,
                values_only,
            );
        }
    }

    Ok(node)
}

/// Expand `SELECT *` into downstream nodes for CTE, derived-table, and raw-table sources.
fn expand_star_into_node(
    node: &mut LineageNode,
    scope: &Scope,
    dialect: Option<DialectType>,
    all_cte_scopes: &[&Scope],
    values_only: bool,
) {
    for (name, source_info) in &scope.sources {
        // CTE source.
        if scope.cte_sources.contains_key(name) {
            if let Some(child_scope) = find_child_scope_in(all_cte_scopes, scope, name) {
                let cte_cols = extract_output_columns(&child_scope.expression);
                if !cte_cols.is_empty() && !cte_cols.contains(&"*".to_string()) {
                    for col in &cte_cols {
                        if let Ok(child) = to_node_inner(
                            ColumnRef::Name(col),
                            child_scope,
                            dialect,
                            "",
                            name,
                            "",
                            false,
                            all_cte_scopes,
                            1,
                            values_only,
                        ) {
                            node.downstream.push(child);
                        }
                    }
                    continue;
                }
            }
        }

        // Derived-table source.
        if source_info.is_scope {
            if let Some(child_scope) = find_child_scope(scope, name) {
                let sub_cols = extract_output_columns(&child_scope.expression);
                if !sub_cols.is_empty() && !sub_cols.contains(&"*".to_string()) {
                    for col in &sub_cols {
                        if let Ok(child) = to_node_inner(
                            ColumnRef::Name(col),
                            child_scope,
                            dialect,
                            "",
                            name,
                            "",
                            false,
                            all_cte_scopes,
                            1,
                            values_only,
                        ) {
                            node.downstream.push(child);
                        }
                    }
                    continue;
                }
            }
        }

        // Raw table — keep as star terminal.
        let child = LineageNode::new(
            format!("{}.*", name),
            Expression::Star(crate::expressions::Star {
                table: None,
                except: None,
                replace: None,
                rename: None,
                trailing_comments: vec![],
            }),
            source_info.expression.clone(),
        );
        node.downstream.push(child);
    }
}

// ---------------------------------------------------------------------------
// Extract output columns from the outermost SELECT
// ---------------------------------------------------------------------------

/// Get the column names / aliases from the outermost SELECT.
/// For set operations (UNION, etc.) walks to the leftmost SELECT.
fn extract_output_columns(expr: &Expression) -> Vec<String> {
    let select = leftmost_select(expr);
    match select {
        Some(Expression::Select(sel)) => sel
            .expressions
            .iter()
            .filter_map(get_alias_or_name)
            .collect(),
        _ => vec![],
    }
}

/// Walk left through set operations to find the leftmost SELECT.
fn leftmost_select(expr: &Expression) -> Option<&Expression> {
    match expr {
        Expression::Select(_) => Some(expr),
        Expression::Union(u) => leftmost_select(&u.left),
        Expression::Intersect(i) => leftmost_select(&i.left),
        Expression::Except(e) => leftmost_select(&e.left),
        Expression::Cte(cte) => leftmost_select(&cte.this),
        _ => None,
    }
}

// ---------------------------------------------------------------------------
// Star resolution — expand `*` through CTE / derived-table scopes
// ---------------------------------------------------------------------------

// ---------------------------------------------------------------------------
// Recursive tree walker — direct port of _process_lineage_for_column()
// ---------------------------------------------------------------------------

/// Recursively walk a `LineageNode` tree and emit `ColumnLineageEdge`s for
/// every terminal (leaf) node.
fn walk_lineage_tree(
    node: &LineageNode,
    target_col: &str,
    lens_type: LensType,
    lens_code: Vec<LensStep>,
    edges: &mut Vec<ColumnLineageEdge>,
) {
    let num_children = node.downstream.len();

    if num_children > 0 {
        // ----- Non-terminal: classify and recurse -----
        let (code_type, new_lens_type) = classify_node(&node.expression, num_children, &lens_type);

        let mut next_lens = lens_code.clone();
        next_lens.push(LensStep {
            expression_sql: node.expression.sql(),
            code_type,
        });

        for child in &node.downstream {
            // When a star node expands to named columns (via CTE/subquery drill),
            // propagate the child's column name as the target instead of "*".
            let child_target = if target_col == "*" {
                let name = extract_column_from_name(&child.name);
                if name != "*" {
                    name
                } else {
                    target_col.to_string()
                }
            } else {
                target_col.to_string()
            };
            walk_lineage_tree(
                child,
                &child_target,
                new_lens_type.clone(),
                next_lens.clone(),
                edges,
            );
        }
        return;
    }

    // ----- Terminal node -----
    //
    // The Rust lineage builder creates terminal nodes with:
    //   expression = Column(table.column)  |  Star  |  Placeholder
    //   source     = Table(TableRef)       |  query expression
    //
    // We identify table terminals by checking node.source for Table
    // (unlike Python sqlglot where the expression itself is the Table).

    match &node.expression {
        Expression::Star(_) => {
            let table_path = extract_star_source(&node.source);
            let source_col = extract_column_from_name(&node.name);

            let mut final_lens = lens_code;
            final_lens.push(LensStep {
                expression_sql: node.expression.sql(),
                code_type: "star".into(),
            });

            edges.push(ColumnLineageEdge {
                source_table: table_path,
                source_column: source_col,
                target_column: target_col.to_string(),
                edge_type: "direct".into(),
                lens_type,
                lens_code: final_lens,
                terminal_kind: TerminalKind::Star,
            });
        }

        Expression::Placeholder(_) => {
            let mut final_lens = lens_code;
            final_lens.push(LensStep {
                expression_sql: node.expression.sql(),
                code_type: "placeholder".into(),
            });

            edges.push(ColumnLineageEdge {
                source_table: String::new(),
                source_column: String::new(),
                target_column: target_col.to_string(),
                edge_type: "direct".into(),
                lens_type,
                lens_code: final_lens,
                terminal_kind: TerminalKind::Placeholder,
            });
        }

        _ => {
            // Check if this is a table terminal (source is a Table expression)
            if let Expression::Table(table_ref) = &node.source {
                let table_path = build_table_path(table_ref);
                let source_col = extract_column_from_name(&node.name);

                let mut final_lens = lens_code;
                final_lens.push(LensStep {
                    expression_sql: node.expression.sql(),
                    code_type: "table".into(),
                });

                edges.push(ColumnLineageEdge {
                    source_table: table_path,
                    source_column: source_col,
                    target_column: target_col.to_string(),
                    edge_type: "direct".into(),
                    lens_type,
                    lens_code: final_lens,
                    terminal_kind: TerminalKind::Table,
                });
            } else if let Expression::Column(col) = &node.expression {
                // Unresolved column reference — emit as unknown
                let table_name = col
                    .table
                    .as_ref()
                    .map(|t| t.name.clone())
                    .unwrap_or_default();

                let mut final_lens = lens_code;
                final_lens.push(LensStep {
                    expression_sql: node.expression.sql(),
                    code_type: "unknown".into(),
                });

                edges.push(ColumnLineageEdge {
                    source_table: table_name,
                    source_column: col.name.name.clone(),
                    target_column: target_col.to_string(),
                    edge_type: "direct".into(),
                    lens_type,
                    lens_code: final_lens,
                    terminal_kind: TerminalKind::Unknown,
                });
            }
            // Other terminal expression types — skip (e.g. COUNT(*) literals)
        }
    }
}

// ---------------------------------------------------------------------------
// Node classification — exact match of Python lines 132-164
// ---------------------------------------------------------------------------

/// Classify a non-terminal node's expression, returning `(code_type, lens_type)`.
///
/// The `code_type` is a string tag for the lens_code step.
/// The `lens_type` may be promoted (Unchanged → Alias or Transformation)
/// but never demoted (Transformation is sticky).
fn classify_node(
    expr: &Expression,
    num_children: usize,
    current_lens_type: &LensType,
) -> (String, LensType) {
    // Determine code_type
    let code_type: String;
    let mut new_lens = current_lens_type.clone();

    if num_children > 1 {
        code_type = "transform".into();
    } else if let Expression::Alias(alias) = expr {
        // Check if Alias wraps a Column with a different name
        if let Expression::Column(col) = &alias.this {
            if col.name.name != alias.alias.name {
                code_type = "alias".into();
            } else {
                code_type = "passthrough".into();
            }
        } else {
            code_type = "transform".into();
        }
    } else {
        // Single child, not an alias — passthrough
        code_type = "passthrough".into();
    };

    // Determine lens_type (only promote, never demote)
    if *current_lens_type != LensType::Transformation {
        if num_children > 1 {
            new_lens = LensType::Transformation;
        } else if let Expression::Alias(alias) = expr {
            if let Expression::Column(col) = &alias.this {
                if col.name.name != alias.alias.name {
                    new_lens = LensType::Alias;
                }
                // same names → keep current
            } else {
                new_lens = LensType::Transformation;
            }
        }
        // else: keep current
    }

    (code_type, new_lens)
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Build an unquoted dot-joined table path: "catalog.schema.name".
fn build_table_path(table: &crate::expressions::TableRef) -> String {
    let mut parts = Vec::with_capacity(3);
    if let Some(ref cat) = table.catalog {
        if !cat.name.is_empty() {
            parts.push(cat.name.as_str());
        }
    }
    if let Some(ref schema) = table.schema {
        if !schema.name.is_empty() {
            parts.push(schema.name.as_str());
        }
    }
    if !table.name.name.is_empty() {
        parts.push(table.name.name.as_str());
    }
    parts.join(".")
}

/// Extract the column name from a lineage node name like "table.column".
/// Falls back to the full string if there's no dot.
fn extract_column_from_name(name: &str) -> String {
    // The node name format is "table.column" — take the last segment
    name.rsplit('.').next().unwrap_or(name).to_string()
}

/// Extract the source table path from a star terminal's source expression.
/// The source is the SELECT containing the star; we pull the table from its FROM clause.
fn extract_star_source(source: &Expression) -> String {
    // Star downstream nodes have source = Table(TableRef) directly
    if let Expression::Table(table_ref) = source {
        return build_table_path(table_ref);
    }
    // Fallback: source may be a Select wrapping the table
    if let Expression::Select(sel) = source {
        if let Some(ref from) = sel.from {
            if let Some(first) = from.expressions.first() {
                if let Expression::Table(table_ref) = first {
                    return build_table_path(table_ref);
                }
            }
        }
    }
    String::new()
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dialects::{Dialect, DialectType};

    fn parse(sql: &str) -> Expression {
        let dialect = Dialect::get(DialectType::Generic);
        let ast = dialect.parse(sql).unwrap();
        ast.into_iter().next().unwrap()
    }

    #[test]
    fn test_simple_lineage() {
        let expr = parse("SELECT a FROM t");
        let node = lineage("a", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "a");
        assert!(!node.downstream.is_empty(), "Should have downstream nodes");
        // Should trace to t.a
        let names = node.downstream_names();
        assert!(
            names.iter().any(|n| n == "t.a"),
            "Expected t.a in downstream, got: {:?}",
            names
        );
    }

    #[test]
    fn test_lineage_walk() {
        let root = LineageNode {
            name: "col_a".to_string(),
            expression: Expression::Null(crate::expressions::Null),
            source: Expression::Null(crate::expressions::Null),
            downstream: vec![LineageNode::new(
                "t.a",
                Expression::Null(crate::expressions::Null),
                Expression::Null(crate::expressions::Null),
            )],
            source_name: String::new(),
            reference_node_name: String::new(),
        };

        let names: Vec<_> = root.walk().map(|n| n.name.clone()).collect();
        assert_eq!(names.len(), 2);
        assert_eq!(names[0], "col_a");
        assert_eq!(names[1], "t.a");
    }

    #[test]
    fn test_aliased_column() {
        let expr = parse("SELECT a + 1 AS b FROM t");
        let node = lineage("b", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "b");
        // Should trace through the expression to t.a
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.iter().any(|n| n.contains("a")),
            "Expected to trace to column a, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_qualified_column() {
        let expr = parse("SELECT t.a FROM t");
        let node = lineage("a", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "a");
        let names = node.downstream_names();
        assert!(
            names.iter().any(|n| n == "t.a"),
            "Expected t.a, got: {:?}",
            names
        );
    }

    #[test]
    fn test_unqualified_column() {
        let expr = parse("SELECT a FROM t");
        let node = lineage("a", &expr, None, false, None).unwrap();

        // Unqualified but single source → resolved to t.a
        let names = node.downstream_names();
        assert!(
            names.iter().any(|n| n == "t.a"),
            "Expected t.a, got: {:?}",
            names
        );
    }

    #[test]
    fn test_lineage_join() {
        let expr = parse("SELECT t.a, s.b FROM t JOIN s ON t.id = s.id");

        let node_a = lineage("a", &expr, None, false, None).unwrap();
        let names_a = node_a.downstream_names();
        assert!(
            names_a.iter().any(|n| n == "t.a"),
            "Expected t.a, got: {:?}",
            names_a
        );

        let node_b = lineage("b", &expr, None, false, None).unwrap();
        let names_b = node_b.downstream_names();
        assert!(
            names_b.iter().any(|n| n == "s.b"),
            "Expected s.b, got: {:?}",
            names_b
        );
    }

    #[test]
    fn test_lineage_derived_table() {
        let expr = parse("SELECT x.a FROM (SELECT a FROM t) AS x");
        let node = lineage("a", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "a");
        // Should trace through the derived table to t.a
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.iter().any(|n| n == "t.a"),
            "Expected to trace through derived table to t.a, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_lineage_cte() {
        let expr = parse("WITH cte AS (SELECT a FROM t) SELECT a FROM cte");
        let node = lineage("a", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "a");
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.iter().any(|n| n == "t.a"),
            "Expected to trace through CTE to t.a, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_lineage_union() {
        let expr = parse("SELECT a FROM t1 UNION SELECT a FROM t2");
        let node = lineage("a", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "a");
        // Should have 2 downstream branches
        assert_eq!(
            node.downstream.len(),
            2,
            "Expected 2 branches for UNION, got {}",
            node.downstream.len()
        );
    }

    #[test]
    fn test_lineage_cte_union() {
        let expr = parse("WITH cte AS (SELECT a FROM t1 UNION SELECT a FROM t2) SELECT a FROM cte");
        let node = lineage("a", &expr, None, false, None).unwrap();

        // Should trace through CTE into both UNION branches
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.len() >= 3,
            "Expected at least 3 nodes for CTE with UNION, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_lineage_star() {
        let expr = parse("SELECT * FROM t");
        let node = lineage("*", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "*");
        // Should have downstream for table t
        assert!(
            !node.downstream.is_empty(),
            "Star should produce downstream nodes"
        );
    }

    #[test]
    fn test_lineage_subquery_in_select() {
        let expr = parse("SELECT (SELECT MAX(b) FROM s) AS x FROM t");
        let node = lineage("x", &expr, None, false, None).unwrap();

        assert_eq!(node.name, "x");
        // Should have traced into the scalar subquery
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.len() >= 2,
            "Expected tracing into scalar subquery, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_lineage_multiple_columns() {
        let expr = parse("SELECT a, b FROM t");

        let node_a = lineage("a", &expr, None, false, None).unwrap();
        let node_b = lineage("b", &expr, None, false, None).unwrap();

        assert_eq!(node_a.name, "a");
        assert_eq!(node_b.name, "b");

        // Each should trace independently
        let names_a = node_a.downstream_names();
        let names_b = node_b.downstream_names();
        assert!(names_a.iter().any(|n| n == "t.a"));
        assert!(names_b.iter().any(|n| n == "t.b"));
    }

    #[test]
    fn test_get_source_tables() {
        let expr = parse("SELECT t.a, s.b FROM t JOIN s ON t.id = s.id");
        let node = lineage("a", &expr, None, false, None).unwrap();

        let tables = get_source_tables(&node);
        assert!(
            tables.contains("t"),
            "Expected source table 't', got: {:?}",
            tables
        );
    }

    #[test]
    fn test_lineage_column_not_found() {
        let expr = parse("SELECT a FROM t");
        let result = lineage("nonexistent", &expr, None, false, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_lineage_nested_cte() {
        let expr = parse(
            "WITH cte1 AS (SELECT a FROM t), \
             cte2 AS (SELECT a FROM cte1) \
             SELECT a FROM cte2",
        );
        let node = lineage("a", &expr, None, false, None).unwrap();

        // Should trace through cte2 → cte1 → t
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.len() >= 3,
            "Expected to trace through nested CTEs, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_trim_selects_true() {
        let expr = parse("SELECT a, b, c FROM t");
        let node = lineage("a", &expr, None, true, None).unwrap();

        // The source should be trimmed to only include 'a'
        if let Expression::Select(select) = &node.source {
            assert_eq!(
                select.expressions.len(),
                1,
                "Trimmed source should have 1 expression, got {}",
                select.expressions.len()
            );
        } else {
            panic!("Expected Select source");
        }
    }

    #[test]
    fn test_trim_selects_false() {
        let expr = parse("SELECT a, b, c FROM t");
        let node = lineage("a", &expr, None, false, None).unwrap();

        // The source should keep all columns
        if let Expression::Select(select) = &node.source {
            assert_eq!(
                select.expressions.len(),
                3,
                "Untrimmed source should have 3 expressions"
            );
        } else {
            panic!("Expected Select source");
        }
    }

    #[test]
    fn test_lineage_expression_in_select() {
        let expr = parse("SELECT a + b AS c FROM t");
        let node = lineage("c", &expr, None, false, None).unwrap();

        // Should trace to both a and b from t
        let all_names: Vec<_> = node.walk().map(|n| n.name.clone()).collect();
        assert!(
            all_names.len() >= 3,
            "Expected to trace a + b to both columns, got: {:?}",
            all_names
        );
    }

    #[test]
    fn test_set_operation_by_index() {
        let expr = parse("SELECT a FROM t1 UNION SELECT b FROM t2");

        // Trace column "a" which is at index 0
        let node = lineage("a", &expr, None, false, None).unwrap();

        // UNION branches should be traced by index
        assert_eq!(node.downstream.len(), 2);
    }

    // -----------------------------------------------------------------------
    // Column lineage tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_col_lineage_simple_passthrough() {
        let r = column_lineage("SELECT a FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty(), "errors: {:?}", r.errors);
        assert_eq!(r.edges.len(), 1);
        let e = &r.edges[0];
        assert_eq!(e.source_table, "t");
        assert_eq!(e.source_column, "a");
        assert_eq!(e.target_column, "a");
        assert_eq!(e.lens_type, LensType::Unchanged);
        assert_eq!(e.terminal_kind, TerminalKind::Table);
    }

    #[test]
    fn test_col_lineage_alias() {
        let r = column_lineage("SELECT a AS b FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 1);
        let e = &r.edges[0];
        assert_eq!(e.source_column, "a");
        assert_eq!(e.target_column, "b");
        assert_eq!(e.lens_type, LensType::Alias);
    }

    #[test]
    fn test_col_lineage_transform() {
        let r = column_lineage("SELECT a + 1 AS c FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 1);
        let e = &r.edges[0];
        assert_eq!(e.source_column, "a");
        assert_eq!(e.target_column, "c");
        assert_eq!(e.lens_type, LensType::Transformation);
    }

    #[test]
    fn test_col_lineage_multi_source_transform() {
        let r = column_lineage("SELECT a + b AS c FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 2, "edges: {:?}", r.edges);
        for e in &r.edges {
            assert_eq!(e.target_column, "c");
            assert_eq!(e.lens_type, LensType::Transformation);
        }
        let cols: HashSet<&str> = r.edges.iter().map(|e| e.source_column.as_str()).collect();
        assert!(cols.contains("a"));
        assert!(cols.contains("b"));
    }

    #[test]
    fn test_col_lineage_join() {
        let r = column_lineage(
            "SELECT t.a, s.b FROM t JOIN s ON t.id = s.id",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 2);
        let edge_a = r.edges.iter().find(|e| e.target_column == "a").unwrap();
        assert_eq!(edge_a.source_table, "t");
        assert_eq!(edge_a.source_column, "a");
        let edge_b = r.edges.iter().find(|e| e.target_column == "b").unwrap();
        assert_eq!(edge_b.source_table, "s");
        assert_eq!(edge_b.source_column, "b");
    }

    #[test]
    fn test_col_lineage_cte() {
        let r = column_lineage(
            "WITH cte AS (SELECT a FROM t) SELECT a FROM cte",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty(), "errors: {:?}", r.errors);
        assert!(!r.edges.is_empty());
        let e = &r.edges[0];
        assert_eq!(e.source_table, "t");
        assert_eq!(e.source_column, "a");
        assert_eq!(e.target_column, "a");
    }

    #[test]
    fn test_col_lineage_derived_table() {
        let r = column_lineage(
            "SELECT x.a FROM (SELECT a FROM t) AS x",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty(), "errors: {:?}", r.errors);
        assert!(!r.edges.is_empty());
        let e = &r.edges[0];
        assert_eq!(e.source_table, "t");
        assert_eq!(e.source_column, "a");
    }

    #[test]
    fn test_col_lineage_union() {
        let r = column_lineage(
            "SELECT a FROM t1 UNION SELECT a FROM t2",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty());
        // Should get edges from both branches
        assert!(r.edges.len() >= 2, "edges: {:?}", r.edges);
    }

    #[test]
    fn test_col_lineage_star() {
        let r = column_lineage("SELECT * FROM t", DialectType::Generic, None);
        // Star may produce Star terminal edges
        assert!(r.errors.is_empty());
        assert!(!r.edges.is_empty());
        let e = &r.edges[0];
        assert_eq!(e.target_column, "*");
    }

    #[test]
    fn test_col_lineage_nested_transform() {
        let r = column_lineage(
            "SELECT UPPER(COALESCE(a, 'default')) AS x FROM t",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty(), "errors: {:?}", r.errors);
        assert_eq!(r.edges.len(), 1);
        let e = &r.edges[0];
        assert_eq!(e.source_column, "a");
        assert_eq!(e.target_column, "x");
        assert_eq!(e.lens_type, LensType::Transformation);
    }

    #[test]
    fn test_col_lineage_case_expression() {
        let r = column_lineage(
            "SELECT CASE WHEN a > 0 THEN b ELSE c END AS result FROM t",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty());
        // CASE with a, b, c should produce multiple edges
        assert!(r.edges.len() >= 2, "edges: {:?}", r.edges);
        for e in &r.edges {
            assert_eq!(e.target_column, "result");
            assert_eq!(e.lens_type, LensType::Transformation);
        }
    }

    #[test]
    fn test_col_lineage_error_resilience() {
        // One valid column and one invalid column — the valid one should still produce edges
        let r = column_lineage("SELECT a FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 1);
    }

    #[test]
    fn test_col_lineage_lens_code_reversed() {
        let r = column_lineage("SELECT a AS b FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 1);
        let e = &r.edges[0];
        // lens_code is reversed to source→target order:
        // first step = table terminal, last step = outermost expression (alias)
        assert!(!e.lens_code.is_empty());
        assert_eq!(
            e.lens_code.first().unwrap().code_type,
            "table",
            "First lens step (after reversal) should be the table terminal"
        );
        assert_eq!(
            e.lens_code.last().unwrap().code_type,
            "alias",
            "Last lens step (after reversal) should be the alias"
        );
    }

    #[test]
    fn test_col_lineage_same_name_alias() {
        // Alias where alias name == column name → Unchanged
        let r = column_lineage("SELECT a AS a FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 1);
        assert_eq!(r.edges[0].lens_type, LensType::Unchanged);
    }

    #[test]
    fn test_col_lineage_qualified_table_path() {
        let r = column_lineage(
            "SELECT a FROM catalog1.schema1.table1",
            DialectType::Generic,
            None,
        );
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 1);
        assert_eq!(r.edges[0].source_table, "catalog1.schema1.table1");
    }

    #[test]
    fn test_col_lineage_multiple_columns() {
        let r = column_lineage("SELECT a, b, c FROM t", DialectType::Generic, None);
        assert!(r.errors.is_empty());
        assert_eq!(r.edges.len(), 3);
        let targets: HashSet<&str> = r.edges.iter().map(|e| e.target_column.as_str()).collect();
        assert!(targets.contains("a"));
        assert!(targets.contains("b"));
        assert!(targets.contains("c"));
    }

    #[test]
    fn test_col_lineage_parse_error() {
        let r = column_lineage("NOT VALID SQL !!!", DialectType::Generic, None);
        assert!(!r.errors.is_empty());
        assert!(r.edges.is_empty());
    }

    #[test]
    fn test_col_lineage_star_through_cte() {
        // Star should expand through the CTE and trace individual columns.
        let sql = "WITH cte AS (SELECT a, b FROM tbl) SELECT * FROM cte";
        let r = column_lineage(sql, DialectType::Snowflake, None);
        eprintln!("edges: {:#?}", r.edges);
        eprintln!("errors: {:#?}", r.errors);
        assert!(r.errors.is_empty(), "errors: {:?}", r.errors);
        // Should get 2 edges (a, b) tracing back to tbl, NOT a single star edge
        assert_eq!(r.edges.len(), 2, "expected 2 edges, got: {:?}", r.edges);
        let src_cols: Vec<&str> = r.edges.iter().map(|e| e.source_column.as_str()).collect();
        assert!(
            src_cols.contains(&"a"),
            "missing source col a, got: {:?}",
            src_cols
        );
        assert!(
            src_cols.contains(&"b"),
            "missing source col b, got: {:?}",
            src_cols
        );
        for e in &r.edges {
            assert_eq!(e.terminal_kind, TerminalKind::Table);
            assert!(
                e.source_table.contains("tbl"),
                "source_table should contain tbl: {:?}",
                e
            );
        }
    }
}
