//! Schema-aware "complete" column lineage.
//!
//! Builds on polyglot-sql's projection lineage to add shared-condition
//! dependencies (WHERE, JOIN ON, GROUP BY, HAVING) — matching the Python
//! `ColumnLineageComplete` semantics.
//!
//! The key insight: in "complete" lineage, columns referenced in shared
//! conditions (WHERE, JOIN ON, GROUP BY, HAVING) are dependencies of
//! **every** output column, because they determine which rows appear.

use std::collections::{BTreeMap, BTreeSet, HashMap};

use datafusion::sql::sqlparser::ast::{
    self, Expr, GroupByExpr, JoinConstraint, JoinOperator, SelectItem, SetExpr, Statement,
};
use datafusion::sql::sqlparser::parser::Parser;
use serde::{Deserialize, Serialize};

use crate::schema::types::SqlDialect;
use crate::validator::dialect::get_parser_dialect;
use polyglot_sql::lineage::{LensType as PolyglotLensType, TerminalKind};

use super::bridge::{edge_to_qualified_ref, normalize_column_name};
use super::column_lineage::polyglot_column_lineage;

/// Whether a source column contributes directly (SELECT projection) or
/// indirectly (WHERE / JOIN / GROUP BY / HAVING).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum LineageType {
    /// Column appears in SELECT projection (direct data flow).
    #[serde(rename = "direct")]
    Direct,
    /// Column appears in a shared condition (WHERE, JOIN ON, GROUP BY, HAVING)
    /// and affects row selection rather than data flow.
    #[serde(rename = "indirect")]
    Indirect,
}

/// How the data is transformed from source to target.
///
/// Serialization matches the Python backend's `LensType` enum values.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum LensClassification {
    /// Column passes through unmodified (polyglot: `Unchanged`).
    #[serde(rename = "Original")]
    Original,
    /// Column was renamed (simple alias of a single column).
    #[serde(rename = "Alias")]
    Alias,
    /// Column was computed via an expression (function, CASE, arithmetic, etc.).
    #[serde(rename = "Transformation")]
    Transformation,
    /// Unable to determine the transformation type.
    #[serde(rename = "Not sure")]
    Unknown,
    /// Shared condition dependency — affects row selection, not data flow.
    #[serde(rename = "Non select")]
    Indirect,
}

/// One step in the expression chain from source to target.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LensStep {
    /// SQL text of the expression at this step.
    pub expression: String,
    /// Classification: "passthrough", "alias", "transform", or terminal kind.
    pub step_type: String,
}

/// A single lineage link from source column to target column.
///
/// A column can appear as both direct AND indirect — e.g. in the SELECT
/// projection AND in a WHERE clause. Each appearance is a separate link
/// in the flat list, so the same (source, target) pair may appear with
/// different `lineage_type` values.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LineageEntry {
    /// Qualified source reference, e.g. `"table1"."col"`.
    pub source: String,
    /// Output column this source feeds into.
    pub target: String,
    /// Direct (SELECT projection) or indirect (WHERE/JOIN/GROUP BY/HAVING).
    pub lineage_type: LineageType,
    /// How the data is transformed.
    pub lens_type: LensClassification,
    /// Expression chain from source to target (empty for indirect edges).
    pub lens_code: Vec<LensStep>,
}

/// Result of complete column lineage analysis.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompleteLineageResult {
    /// Map from output column name to sorted list of source references.
    /// Each source reference is a qualified `"table"."column"` string.
    /// Backward-compatible flat format (includes both direct and indirect).
    pub column_dict: HashMap<String, Vec<String>>,
    /// Flat list of lineage links with full metadata.
    /// Each link carries source, target, lineage_type, lens_type, and
    /// expression chain. Indirect sources appear once per output column.
    pub column_lineage: Vec<LineageEntry>,
    /// Source tables referenced by the query.
    pub source_tables: Vec<String>,
    /// Errors encountered during lineage tracing.
    pub errors: Vec<String>,
    /// Default database applied during qualification (if any).
    pub default_database: Option<String>,
    /// Default schema applied during qualification (if any).
    pub default_schema: Option<String>,
}

/// Options for column lineage analysis.
///
/// All fields are optional with sensible defaults. Pass `None` to
/// `column_lineage()` for projection-only lineage with no qualification.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ColumnLineageOptions {
    /// Default database for qualifying unqualified table references.
    pub default_database: Option<String>,
    /// Default schema for qualifying unqualified table references.
    pub default_schema: Option<String>,
    /// When true, inject shared-condition dependencies (WHERE, JOIN ON,
    /// GROUP BY, HAVING) as indirect lineage edges on every output column.
    /// Default: false (projection-only, matching Snowflake ACCESS_USAGE
    /// "direct" lineage semantics).
    pub inject_shared: bool,
}

/// Convert polyglot-sql's `LensType` to our `LensClassification`.
fn polyglot_lens_to_classification(lt: &PolyglotLensType) -> LensClassification {
    match lt {
        PolyglotLensType::Unchanged => LensClassification::Original,
        PolyglotLensType::Alias => LensClassification::Alias,
        PolyglotLensType::Transformation => LensClassification::Transformation,
    }
}

/// Convert polyglot-sql's `LensStep` list to our `LensStep` list.
fn polyglot_lens_code_to_steps(steps: &[polyglot_sql::lineage::LensStep]) -> Vec<LensStep> {
    steps
        .iter()
        .map(|s| LensStep {
            expression: s.expression_sql.clone(),
            step_type: s.code_type.clone(),
        })
        .collect()
}

// ===========================================================================
// Schema flattening
// ===========================================================================

/// Flattened schema: table name (possibly multi-part like "db.schema.table1")
/// → list of column names.
type FlatSchema = BTreeMap<String, Vec<String>>;

/// CTE column-level lineage dictionary.
/// Maps CTE name → { output column name → set of base table source refs }.
/// Built in CTE declaration order so later CTEs can reference earlier ones.
type CteDict = BTreeMap<String, BTreeMap<String, BTreeSet<String>>>;

/// Flatten a potentially nested schema into a table → columns map.
///
/// Handles both flat schemas like `{"table1": {"a": "int"}}` and nested
/// schemas like `{"db": {"schema": {"table1": {"a": "int"}}}}`.
fn flatten_schema(schema: &HashMap<String, serde_json::Value>) -> FlatSchema {
    let mut result = BTreeMap::new();
    for (key, val) in schema {
        flatten_schema_recursive(val, &[key.as_str()], &mut result);
    }
    result
}

fn flatten_schema_recursive(val: &serde_json::Value, path: &[&str], result: &mut FlatSchema) {
    if let serde_json::Value::Object(obj) = val {
        let has_string_values = obj.values().any(|v| v.is_string());
        let has_object_values = obj.values().any(|v| v.is_object());

        if has_string_values && !has_object_values {
            let table_name = path.join(".");
            let columns: Vec<String> = obj.keys().cloned().collect();
            result.insert(table_name, columns);
        } else if has_object_values {
            for (key, child) in obj {
                let mut new_path = path.to_vec();
                new_path.push(key.as_str());
                flatten_schema_recursive(child, &new_path, result);
            }
        }
    }
}

/// Check if a flat_schema table has a column (case-insensitive).
/// Uses `find_schema_table` for multi-part name resolution.
fn schema_has_column(flat_schema: &FlatSchema, table: &str, column: &str) -> bool {
    let col_lower = column.to_lowercase();
    if let Some(canonical) = find_schema_table(flat_schema, table)
        && let Some(cols) = flat_schema.get(canonical)
    {
        return cols.iter().any(|c| c.to_lowercase() == col_lower);
    }
    false
}

/// Normalize over-qualified source references by stripping database/schema
/// prefixes when the flat schema only has simple table names.
///
/// Polyglot-sql can produce `"DW"."MEMBER_PROFILE"."COL"` from SQL like
/// `FROM "DW"."MEMBER_PROFILE"`. When the schema key is just `MEMBER_PROFILE`
/// (no dots), this rewrites to `"MEMBER_PROFILE"."COL"`.
///
/// Skips normalization when schema keys themselves contain dots (nested schemas
/// like `db.schema.table`), preserving the full path for those cases.
fn normalize_source_ref_prefixes(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) {
    // If any schema key contains a dot, schema is nested — don't strip.
    if flat_schema.keys().any(|k| k.contains('.')) {
        return;
    }

    for sources in column_dict.values_mut() {
        let normalized: BTreeSet<String> = sources
            .iter()
            .map(|source| {
                // Count quoted parts: "A"."B"."C" has 3 parts
                let parts: Vec<&str> = source.split("\".\"").collect();
                if parts.len() <= 2 {
                    // Already 2-part or less — no stripping needed
                    return source.clone();
                }
                // Strip to 2-part: keep only last two segments (table + column).
                // "DB"."SCHEMA"."TABLE"."COL" → "TABLE"."COL"
                let table_part = parts[parts.len() - 2].trim_matches('"');
                let col_part = parts[parts.len() - 1].trim_matches('"');
                format!("\"{}\".\"{}\"", table_part, col_part)
            })
            .map(|s| match dialect {
                SqlDialect::Snowflake => s.to_uppercase(),
                _ => s,
            })
            .collect();
        *sources = normalized;
    }
}

/// Find a table name in the schema (case-insensitive match).
/// Handles multi-part names like `db.schema.table` or `schema.table` by
/// matching the last segment against schema keys when no exact match is found.
fn find_schema_table<'a>(flat_schema: &'a FlatSchema, name: &str) -> Option<&'a str> {
    let lower = name.to_lowercase();
    // Exact match first
    if let Some(k) = flat_schema.keys().find(|k| k.to_lowercase() == lower) {
        return Some(k.as_str());
    }
    // Strip quotes and try again
    let stripped = lower
        .trim_matches('"')
        .replace("\".", ".")
        .replace(".\"", ".");
    if stripped != lower
        && let Some(k) = flat_schema.keys().find(|k| k.to_lowercase() == stripped)
    {
        return Some(k.as_str());
    }
    // Multi-part fallback: match the last segment (table name only)
    if let Some(dot_pos) = name.rfind('.') {
        let table_part = name[dot_pos + 1..].trim_matches('"').to_lowercase();
        if !table_part.is_empty() {
            return flat_schema
                .keys()
                .find(|k| k.trim_matches('"').to_lowercase() == table_part)
                .map(|k| k.as_str());
        }
    }
    None
}

/// Find a CTE name case-insensitively.
fn find_cte_name<'a>(cte_names: &'a BTreeSet<String>, name: &str) -> Option<&'a String> {
    cte_names.iter().find(|c| c.eq_ignore_ascii_case(name))
}

/// Find a subquery source case-insensitively.
fn find_subquery_source<'a>(
    sources: &'a BTreeMap<String, Vec<String>>,
    name: &str,
) -> Option<&'a String> {
    sources.keys().find(|k| k.eq_ignore_ascii_case(name))
}

/// Find an alias case-insensitively, returning the mapped value.
fn find_alias<'a>(alias_map: &'a AliasMap, name: &str) -> Option<&'a String> {
    alias_map
        .iter()
        .find(|(k, _)| k.eq_ignore_ascii_case(name))
        .map(|(_, v)| v)
}

/// Find a CTE in the dict (case-insensitive).
fn find_cte_in_dict<'a>(
    cte_dict: &'a CteDict,
    name: &str,
) -> Option<&'a BTreeMap<String, BTreeSet<String>>> {
    cte_dict.get(name).or_else(|| {
        let lower = name.to_lowercase();
        cte_dict
            .iter()
            .find(|(k, _)| k.to_lowercase() == lower)
            .map(|(_, v)| v)
    })
}

/// Find schema tables that have a given column (case-insensitive).
/// Returns the canonical table names.
fn find_tables_with_column<'a>(flat_schema: &'a FlatSchema, column: &str) -> Vec<&'a str> {
    let col_lower = column.to_lowercase();
    flat_schema
        .iter()
        .filter(|(_, cols)| cols.iter().any(|c| c.to_lowercase() == col_lower))
        .map(|(t, _)| t.as_str())
        .collect()
}

// ===========================================================================
// Comment stripping
// ===========================================================================

fn strip_comments(sql: &str) -> String {
    let mut result = String::with_capacity(sql.len());
    let chars: Vec<char> = sql.chars().collect();
    let len = chars.len();
    let mut i = 0;

    while i < len {
        if i + 1 < len && chars[i] == '-' && chars[i + 1] == '-' {
            while i < len && chars[i] != '\n' {
                i += 1;
            }
        } else if i + 1 < len && chars[i] == '/' && chars[i + 1] == '*' {
            i += 2;
            while i + 1 < len && !(chars[i] == '*' && chars[i + 1] == '/') {
                i += 1;
            }
            if i + 1 < len {
                i += 2;
            }
        } else {
            result.push(chars[i]);
            i += 1;
        }
    }

    result
}

// ===========================================================================
// Query structural analysis
// ===========================================================================

/// Table alias mapping: alias → actual table name (possibly multi-part).
type AliasMap = BTreeMap<String, String>;

/// A raw column reference extracted from the AST.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct RawColumnRef {
    table: Option<String>,
    column: String,
}

/// Structural analysis of a parsed SQL query.
struct QueryInfo {
    /// Alias map for the outermost (or leftmost) SELECT.
    alias_map: AliasMap,
    /// CTE names.
    cte_names: BTreeSet<String>,
    /// CTE name → list of base tables it references.
    cte_source_tables: BTreeMap<String, Vec<String>>,
    /// CTEs actually referenced in the main query (directly or transitively).
    used_ctes: BTreeSet<String>,
    /// Whether the outermost query is a set operation.
    is_set_operation: bool,
    /// Tables on the left side of a set operation (including CTE names).
    left_side_tables: BTreeSet<String>,
    /// Parsed SELECT items from the leftmost SELECT.
    select_items: Vec<ast::SelectItem>,
    /// The parsed query.
    query: Box<ast::Query>,
    /// Subquery alias → source tables within the subquery.
    subquery_source_tables: BTreeMap<String, Vec<String>>,
    /// CTE name → output column names (from its SELECT list).
    cte_output_columns: BTreeMap<String, Vec<String>>,
    /// CTE name → its body's alias map.
    cte_alias_maps: BTreeMap<String, AliasMap>,
    /// For INSERT INTO: target column names (from explicit column list or schema).
    /// Empty for non-INSERT queries.
    insert_target_columns: Vec<String>,
    /// LATERAL FLATTEN alias → input column references.
    /// Maps virtual table aliases (e.g., `F` from `LATERAL FLATTEN(input => col) F`)
    /// to the columns feeding the FLATTEN. `F.VALUE` should trace back to those columns
    /// rather than appearing as a raw source.
    lateral_flatten_inputs: BTreeMap<String, Vec<RawColumnRef>>,
}

fn analyze_query(stmts: &[Statement]) -> Option<QueryInfo> {
    for stmt in stmts {
        let (query, insert_target_columns) = match stmt {
            Statement::Query(query) => (query.clone(), vec![]),
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    // Extract INSERT INTO column names for positional mapping
                    let target_cols: Vec<String> =
                        ins.columns.iter().map(|c| c.value.clone()).collect();
                    (src.clone(), target_cols)
                } else {
                    continue;
                }
            }
            _ => continue,
        };

        let cte_names = collect_cte_names(&query);
        let cte_source_tables = build_cte_source_tables(&query, &cte_names);
        let mut subquery_source_tables = BTreeMap::new();
        let alias_map = build_alias_map_from_body(&query.body, &mut subquery_source_tables);

        // CTE output columns and alias maps
        let mut cte_output_columns = BTreeMap::new();
        let mut cte_alias_maps = BTreeMap::new();
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                let name = cte.alias.name.value.clone();
                let cols = extract_select_column_names(&cte.query.body);
                cte_output_columns.insert(name.clone(), cols);
                let mut sub = BTreeMap::new();
                let am = build_alias_map_from_body(&cte.query.body, &mut sub);
                subquery_source_tables.extend(sub);
                cte_alias_maps.insert(name, am);
            }
        }

        let used_ctes = find_used_ctes(&query, &cte_names);
        let (is_set_operation, left_side_tables) = detect_set_operation(&query.body);
        let select_items = extract_leftmost_select_items(&query.body);
        let lateral_flatten_inputs = collect_flatten_inputs(&query);
        return Some(QueryInfo {
            alias_map,
            cte_names,
            cte_source_tables,
            used_ctes,
            is_set_operation,
            left_side_tables,
            select_items,
            query,
            subquery_source_tables,
            cte_output_columns,
            cte_alias_maps,
            insert_target_columns,
            lateral_flatten_inputs,
        });
    }
    None
}

fn collect_cte_names(query: &ast::Query) -> BTreeSet<String> {
    let mut names = BTreeSet::new();
    if let Some(with) = &query.with {
        for cte in &with.cte_tables {
            names.insert(cte.alias.name.value.clone());
        }
    }
    names
}

/// Find which CTEs are actually used (referenced from the main query body or
/// transitively from other used CTEs).
fn find_used_ctes(query: &ast::Query, cte_names: &BTreeSet<String>) -> BTreeSet<String> {
    let mut used = BTreeSet::new();

    // Collect table names referenced in the main query body
    let main_tables = extract_table_names_from_body(&query.body);
    for t in &main_tables {
        if cte_names.contains(t) {
            used.insert(t.clone());
        }
    }

    // Transitively: used CTE references another CTE → that CTE is used too
    let mut changed = true;
    while changed {
        changed = false;
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                let name = cte.alias.name.value.clone();
                if used.contains(&name) {
                    let inner = extract_table_names_from_body(&cte.query.body);
                    for t in &inner {
                        if cte_names.contains(t) && !used.contains(t) {
                            used.insert(t.clone());
                            changed = true;
                        }
                    }
                }
            }
        }
    }

    used
}

/// Build a map of CTE name → list of base tables it references (transitively).
fn build_cte_source_tables(
    query: &ast::Query,
    cte_names: &BTreeSet<String>,
) -> BTreeMap<String, Vec<String>> {
    let mut map = BTreeMap::new();
    if let Some(with) = &query.with {
        for cte in &with.cte_tables {
            let name = cte.alias.name.value.clone();
            let tables = extract_table_names_from_body(&cte.query.body);
            map.insert(name, tables);
        }
        // Resolve transitive CTE references
        let cte_order: Vec<String> = with
            .cte_tables
            .iter()
            .map(|c| c.alias.name.value.clone())
            .collect();
        for cte_name in &cte_order {
            let tables = map.get(cte_name).cloned().unwrap_or_default();
            let mut resolved = Vec::new();
            for t in &tables {
                if cte_names.contains(t) {
                    if let Some(inner) = map.get(t) {
                        resolved.extend(inner.clone());
                    }
                } else {
                    resolved.push(t.clone());
                }
            }
            resolved.sort();
            resolved.dedup();
            map.insert(cte_name.clone(), resolved);
        }
    }
    map
}

fn extract_table_names_from_body(body: &SetExpr) -> Vec<String> {
    match body {
        SetExpr::Select(sel) => {
            let mut tables = Vec::new();
            for twj in &sel.from {
                extract_table_name_from_factor(&twj.relation, &mut tables);
                for join in &twj.joins {
                    extract_table_name_from_factor(&join.relation, &mut tables);
                }
            }
            tables
        }
        SetExpr::SetOperation { left, right, .. } => {
            let mut tables = extract_table_names_from_body(left);
            tables.extend(extract_table_names_from_body(right));
            tables
        }
        SetExpr::Query(q) => extract_table_names_from_body(&q.body),
        _ => vec![],
    }
}

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

/// Collect LATERAL FLATTEN aliases and their input column references.
///
/// Maps each virtual alias (e.g., `F` from `LATERAL FLATTEN(input => col) F`)
/// to the column references that feed the FLATTEN. This allows tracing
/// `F.VALUE` back to the actual source column being flattened.
fn collect_flatten_inputs(query: &ast::Query) -> BTreeMap<String, Vec<RawColumnRef>> {
    let mut map = BTreeMap::new();
    collect_flatten_inputs_from_body(&query.body, &mut map);
    if let Some(with) = &query.with {
        for cte in &with.cte_tables {
            collect_flatten_inputs_from_body(&cte.query.body, &mut map);
        }
    }
    map
}

fn collect_flatten_inputs_from_body(body: &SetExpr, map: &mut BTreeMap<String, Vec<RawColumnRef>>) {
    match body {
        SetExpr::Select(sel) => {
            for twj in &sel.from {
                collect_flatten_inputs_from_factor(&twj.relation, map);
                for join in &twj.joins {
                    collect_flatten_inputs_from_factor(&join.relation, map);
                }
            }
        }
        SetExpr::SetOperation { left, right, .. } => {
            collect_flatten_inputs_from_body(left, map);
            collect_flatten_inputs_from_body(right, map);
        }
        SetExpr::Query(q) => collect_flatten_inputs_from_body(&q.body, map),
        _ => {}
    }
}

fn collect_flatten_inputs_from_factor(
    factor: &ast::TableFactor,
    map: &mut BTreeMap<String, Vec<RawColumnRef>>,
) {
    match factor {
        ast::TableFactor::Function {
            alias: Some(a),
            args,
            ..
        } => {
            let mut refs = Vec::new();
            for arg in args {
                let expr = match arg {
                    ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e))
                    | ast::FunctionArg::Named {
                        arg: ast::FunctionArgExpr::Expr(e),
                        ..
                    }
                    | ast::FunctionArg::ExprNamed {
                        arg: ast::FunctionArgExpr::Expr(e),
                        ..
                    } => e,
                    _ => continue,
                };
                extract_columns_from_expr(expr, &mut refs);
            }
            map.insert(a.name.value.clone(), refs);
        }
        ast::TableFactor::Derived { subquery, .. } => {
            collect_flatten_inputs_from_body(&subquery.body, map);
        }
        _ => {}
    }
}

/// Replace LATERAL FLATTEN alias refs in column_dict with resolved input columns.
///
/// Scans all source refs in `column_dict`. If a ref like `"F"."VALUE"` matches
/// a known FLATTEN alias, it is replaced with the resolved input column(s).
fn resolve_flatten_refs_in_dict(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    flatten_inputs: &BTreeMap<String, Vec<RawColumnRef>>,
    alias_map: &AliasMap,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    for sources in column_dict.values_mut() {
        let mut to_remove = Vec::new();
        let mut to_add = Vec::new();

        for source in sources.iter() {
            // Parse the source ref to extract the table part
            // Format: "TABLE"."COL" or just TABLE.COL
            let stripped = source.replace('"', "");
            let parts: Vec<&str> = stripped.split('.').collect();
            if parts.len() < 2 {
                continue;
            }
            let table_part = parts[parts.len() - 2];

            // Check if this table is a FLATTEN alias
            if let Some(input_refs) = flatten_inputs
                .iter()
                .find(|(a, _)| a.eq_ignore_ascii_case(table_part))
                .map(|(_, refs)| refs)
            {
                to_remove.push(source.clone());
                for raw in input_refs {
                    let resolved = resolve_ref_with_cte_dict(
                        raw,
                        alias_map,
                        flat_schema,
                        cte_dict,
                        dialect,
                        None,
                    );
                    to_add.extend(resolved);
                }
            }
        }

        for r in &to_remove {
            sources.remove(r);
        }
        for r in to_add {
            sources.insert(r);
        }
    }
}

/// Build alias map from a query body, tracking subquery source tables.
fn build_alias_map_from_body(
    body: &SetExpr,
    subquery_sources: &mut BTreeMap<String, Vec<String>>,
) -> AliasMap {
    match body {
        SetExpr::Select(sel) => build_alias_map(&sel.from, subquery_sources),
        SetExpr::SetOperation { left, .. } => build_alias_map_from_body(left, subquery_sources),
        SetExpr::Query(q) => build_alias_map_from_body(&q.body, subquery_sources),
        _ => AliasMap::new(),
    }
}

fn build_alias_map(
    from: &[ast::TableWithJoins],
    subquery_sources: &mut BTreeMap<String, Vec<String>>,
) -> AliasMap {
    let mut aliases = AliasMap::new();
    for twj in from {
        collect_table_aliases(&twj.relation, &mut aliases, subquery_sources);
        for join in &twj.joins {
            collect_table_aliases(&join.relation, &mut aliases, subquery_sources);
        }
    }
    aliases
}

fn collect_table_aliases(
    factor: &ast::TableFactor,
    aliases: &mut AliasMap,
    subquery_sources: &mut BTreeMap<String, Vec<String>>,
) {
    match factor {
        ast::TableFactor::Table { name, alias, .. } => {
            let table_name = name.to_string();
            aliases.insert(table_name.clone(), table_name.clone());
            if let Some(a) = alias {
                aliases.insert(a.name.value.clone(), table_name);
            }
        }
        ast::TableFactor::Derived {
            alias, subquery, ..
        } => {
            // Track what tables this subquery references
            let inner_tables = extract_table_names_from_body(&subquery.body);
            if let Some(a) = alias {
                let alias_name = a.name.value.clone();
                subquery_sources.insert(alias_name.clone(), inner_tables.clone());
                aliases.insert(alias_name, "__subquery__".to_string());
            } else {
                // Un-aliased derived subquery: register under synthetic name
                // so its output columns are reachable via CTE dict.
                let synthetic = "__derived__".to_string();
                subquery_sources.insert(synthetic.clone(), inner_tables.clone());
                aliases.insert(synthetic, "__subquery__".to_string());
            }
            // Also add inner table aliases so they're available for resolution,
            // but ONLY when the subquery has an explicit alias (scoped access).
            // Without an alias, inner tables must NOT leak into the outer scope —
            // the outer SELECT should resolve through the subquery's output columns.
            if alias.is_some()
                && let SetExpr::Select(sel) = subquery.body.as_ref()
            {
                let mut sub = BTreeMap::new();
                let inner_aliases = build_alias_map(&sel.from, &mut sub);
                subquery_sources.extend(sub);
                for table in inner_aliases.values() {
                    if !table.starts_with("__subquery__") {
                        aliases.insert(table.clone(), table.clone());
                    }
                }
            }
        }
        _ => {}
    }
}

/// Detect if the query body is a set operation and extract left-side table names.
fn detect_set_operation(body: &SetExpr) -> (bool, BTreeSet<String>) {
    match body {
        SetExpr::SetOperation { left, .. } => {
            let left_tables: BTreeSet<String> =
                extract_table_names_from_body(left).into_iter().collect();
            (true, left_tables)
        }
        _ => (false, BTreeSet::new()),
    }
}

/// Extract SELECT items from the leftmost SELECT in a query body.
fn extract_leftmost_select_items(body: &SetExpr) -> Vec<ast::SelectItem> {
    match body {
        SetExpr::Select(sel) => sel.projection.clone(),
        SetExpr::SetOperation { left, .. } => extract_leftmost_select_items(left),
        SetExpr::Query(q) => extract_leftmost_select_items(&q.body),
        _ => vec![],
    }
}

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

// ===========================================================================
// CTE dict: pre-computed column-level lineage for each CTE
//
// Modeled after Python's _run_cte_lineage: processes CTEs in declaration
// order, building a complete {column → [base table refs]} mapping for each.
// Later CTEs resolve through earlier CTEs transitively.
// ===========================================================================

/// Build the CTE dict by processing each CTE's SELECT list in declaration order.
fn build_cte_dict(query: &ast::Query, flat_schema: &FlatSchema, dialect: SqlDialect) -> CteDict {
    let mut cte_dict = CteDict::new();

    let with = match &query.with {
        Some(w) => w,
        None => return cte_dict,
    };

    for cte in &with.cte_tables {
        let name = cte.alias.name.value.clone();
        let mut col_map: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();

        // Process the CTE body (handles UNION by merging all branches)
        process_body_for_cte_dict(
            &cte.query.body,
            flat_schema,
            &mut cte_dict,
            dialect,
            &mut col_map,
        );

        // Store with both original and uppercase keys for case-insensitive lookup
        let name_upper = name.to_uppercase();
        if name != name_upper {
            cte_dict.insert(name_upper, col_map.clone());
        }
        cte_dict.insert(name, col_map);
    }

    cte_dict
}

/// Process a SetExpr body for CTE dict building (handles UNION/INTERSECT/EXCEPT).
fn process_body_for_cte_dict(
    body: &SetExpr,
    flat_schema: &FlatSchema,
    cte_dict: &mut CteDict,
    dialect: SqlDialect,
    col_map: &mut BTreeMap<String, BTreeSet<String>>,
) {
    match body {
        SetExpr::Select(sel) => {
            let mut sub = BTreeMap::new();
            let alias_map = build_alias_map(&sel.from, &mut sub);

            // Build dict entries for inline derived subqueries BEFORE processing
            // select items, so references like BUTTON_ORDERS from an inner
            // (SELECT ... FROM base_table) are resolvable during CTE dict building.
            for twj in &sel.from {
                build_inline_subquery_dict(&twj.relation, flat_schema, cte_dict, dialect);
                for join in &twj.joins {
                    build_inline_subquery_dict(&join.relation, flat_schema, cte_dict, dialect);
                }
            }

            for item in &sel.projection {
                process_select_item_for_cte_dict(
                    item,
                    &alias_map,
                    flat_schema,
                    cte_dict,
                    dialect,
                    col_map,
                );
            }
        }
        SetExpr::SetOperation { left, right, .. } => {
            // Process both branches and merge (UNION sources from both sides)
            process_body_for_cte_dict(left, flat_schema, cte_dict, dialect, col_map);
            process_body_for_cte_dict(right, flat_schema, cte_dict, dialect, col_map);
        }
        SetExpr::Query(q) => {
            process_body_for_cte_dict(&q.body, flat_schema, cte_dict, dialect, col_map);
        }
        _ => {}
    }
}

/// Build dict entries for derived subqueries in a FROM clause factor.
/// Called BEFORE processing SELECT items so inline subquery columns are resolvable.
fn build_inline_subquery_dict(
    factor: &ast::TableFactor,
    flat_schema: &FlatSchema,
    dict: &mut CteDict,
    dialect: SqlDialect,
) {
    match factor {
        ast::TableFactor::Derived {
            alias, subquery, ..
        } => {
            // Recurse into nested subqueries first (bottom-up)
            if let SetExpr::Select(inner_sel) = subquery.body.as_ref() {
                for twj in &inner_sel.from {
                    build_inline_subquery_dict(&twj.relation, flat_schema, dict, dialect);
                    for join in &twj.joins {
                        build_inline_subquery_dict(&join.relation, flat_schema, dict, dialect);
                    }
                }
            }
            // Build dict entry for this subquery
            let alias_name = if let Some(a) = alias {
                a.name.value.clone()
            } else {
                "__derived__".to_string()
            };
            if find_cte_in_dict(dict, &alias_name).is_none() {
                let mut sub_col_map = BTreeMap::new();
                process_body_for_cte_dict(
                    &subquery.body,
                    flat_schema,
                    dict,
                    dialect,
                    &mut sub_col_map,
                );
                if !sub_col_map.is_empty() {
                    let name_upper = alias_name.to_uppercase();
                    if alias_name != name_upper {
                        dict.insert(name_upper, sub_col_map.clone());
                    }
                    dict.insert(alias_name, sub_col_map);
                }
            }
        }
        ast::TableFactor::NestedJoin {
            table_with_joins, ..
        } => {
            build_inline_subquery_dict(&table_with_joins.relation, flat_schema, dict, dialect);
            for join in &table_with_joins.joins {
                build_inline_subquery_dict(&join.relation, flat_schema, dict, dialect);
            }
        }
        _ => {}
    }
}

/// Process a single SELECT item for CTE dict.
fn process_select_item_for_cte_dict(
    item: &SelectItem,
    alias_map: &AliasMap,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    col_map: &mut BTreeMap<String, BTreeSet<String>>,
) {
    match item {
        SelectItem::Wildcard(_) => {
            // Expand * from each source table in scope
            for (alias_key, actual_table) in alias_map.iter() {
                if actual_table == "__subquery__" {
                    // Subquery alias — expand using the alias name (may be in dict)
                    expand_table_into_cte_dict(alias_key, flat_schema, cte_dict, dialect, col_map);
                    continue;
                }
                expand_table_into_cte_dict(actual_table, flat_schema, cte_dict, dialect, col_map);
            }
        }
        SelectItem::QualifiedWildcard(name, _) => {
            let prefix = name.to_string();
            let (alias_key, actual_table) = alias_map
                .iter()
                .find(|(k, _)| k.eq_ignore_ascii_case(&prefix))
                .map(|(k, v)| (k.as_str(), v.as_str()))
                .unwrap_or((prefix.as_str(), prefix.as_str()));
            if actual_table == "__subquery__" {
                expand_table_into_cte_dict(alias_key, flat_schema, cte_dict, dialect, col_map);
            } else {
                expand_table_into_cte_dict(actual_table, flat_schema, cte_dict, dialect, col_map);
            }
        }
        SelectItem::ExprWithAlias { expr, alias } => {
            let col_name = normalize_column_name(&alias.value, dialect);
            let mut sources = BTreeSet::new();
            resolve_expr_for_cte_dict(
                expr,
                alias_map,
                flat_schema,
                cte_dict,
                dialect,
                &mut sources,
                Some(col_map),
            );
            col_map.entry(col_name).or_default().extend(sources);
        }
        SelectItem::UnnamedExpr(expr) => {
            let col_name = match expr {
                Expr::Identifier(ident) => normalize_column_name(&ident.value, dialect),
                Expr::CompoundIdentifier(parts) => {
                    if let Some(last) = parts.last() {
                        normalize_column_name(&last.value, dialect)
                    } else {
                        return;
                    }
                }
                // For complex expressions (functions, CASE, CAST, etc.) without
                // an explicit alias: derive a column name from the expression text.
                // SQL engines assign implicit names like "CONCAT(A,B)" or "CASE".
                // This ensures the CTE dict captures ALL output columns, enabling
                // outer queries to trace through via SELECT * expansion.
                _ => {
                    let expr_str = expr.to_string();
                    normalize_column_name(&expr_str, dialect)
                }
            };
            let mut sources = BTreeSet::new();
            resolve_expr_for_cte_dict(
                expr,
                alias_map,
                flat_schema,
                cte_dict,
                dialect,
                &mut sources,
                Some(col_map),
            );
            col_map.entry(col_name).or_default().extend(sources);
        }
    }
}

// ===========================================================================
// Subquery dict — extend the CTE dict with derived subquery alias mappings
// ===========================================================================

/// Recursively process a CTE's full query: its own inner WITH clause (if any)
/// followed by its body. This ensures nested CTE definitions like
/// `WITH outer AS (WITH inner AS (...) SELECT ...)` correctly build dict entries
/// for `inner` before processing `outer`'s body.
fn extend_subquery_from_cte_query(
    cte_query: &ast::Query,
    flat_schema: &FlatSchema,
    dict: &mut CteDict,
    dialect: SqlDialect,
) {
    // Process nested WITH clause within this CTE's own query
    if let Some(with) = &cte_query.with {
        for cte in &with.cte_tables {
            let name = cte.alias.name.value.clone();
            // Recursively handle deeper nesting
            extend_subquery_from_cte_query(&cte.query, flat_schema, dict, dialect);
            // Build dict entry for this nested CTE
            if find_cte_in_dict(dict, &name).is_none() {
                let mut col_map = BTreeMap::new();
                process_body_for_cte_dict(
                    &cte.query.body,
                    flat_schema,
                    dict,
                    dialect,
                    &mut col_map,
                );
                if !col_map.is_empty() {
                    let name_upper = name.to_uppercase();
                    if name != name_upper {
                        dict.insert(name_upper, col_map.clone());
                    }
                    dict.insert(name, col_map);
                }
            }
        }
    }
    // Then process the body (derived subqueries in FROM, etc.)
    extend_subquery_from_body(&cte_query.body, flat_schema, dict, dialect);
}

/// Walk the AST and add column→source mappings for every derived subquery alias.
/// Processes bottom-up (innermost subqueries first) so outer subqueries can
/// reference inner ones through the dict.
fn extend_with_subquery_dict(
    query: &ast::Query,
    flat_schema: &FlatSchema,
    dict: &mut CteDict,
    dialect: SqlDialect,
) {
    // Walk main query body
    extend_subquery_from_body(&query.body, flat_schema, dict, dialect);
    // Walk CTE bodies (CTEs may contain derived subqueries)
    if let Some(with) = &query.with {
        for cte in &with.cte_tables {
            extend_subquery_from_body(&cte.query.body, flat_schema, dict, dialect);
        }
    }
}

fn extend_subquery_from_body(
    body: &SetExpr,
    flat_schema: &FlatSchema,
    dict: &mut CteDict,
    dialect: SqlDialect,
) {
    match body {
        SetExpr::Select(sel) => {
            for twj in &sel.from {
                extend_subquery_from_factor(&twj.relation, flat_schema, dict, dialect);
                for join in &twj.joins {
                    extend_subquery_from_factor(&join.relation, flat_schema, dict, dialect);
                }
            }
        }
        SetExpr::SetOperation { left, right, .. } => {
            extend_subquery_from_body(left, flat_schema, dict, dialect);
            extend_subquery_from_body(right, flat_schema, dict, dialect);
        }
        SetExpr::Query(q) => {
            // Handle inner WITH clause and body via recursive helper
            extend_subquery_from_cte_query(q, flat_schema, dict, dialect);
        }
        _ => {}
    }
}

fn extend_subquery_from_factor(
    factor: &ast::TableFactor,
    flat_schema: &FlatSchema,
    dict: &mut CteDict,
    dialect: SqlDialect,
) {
    match factor {
        ast::TableFactor::Derived {
            alias, subquery, ..
        } => {
            // Handle inner WITH clause and body via recursive helper
            // (handles nested CTEs like `WITH outer AS (WITH inner AS (...) ...)`)
            extend_subquery_from_cte_query(subquery, flat_schema, dict, dialect);

            // Build column map for this subquery alias (or synthetic name if un-aliased)
            let alias_name = if let Some(a) = alias {
                a.name.value.clone()
            } else {
                "__derived__".to_string()
            };
            // Only build if not already in dict (avoid overwriting CTE entries)
            if find_cte_in_dict(dict, &alias_name).is_none() {
                let mut col_map = BTreeMap::new();
                process_body_for_cte_dict(&subquery.body, flat_schema, dict, dialect, &mut col_map);
                if !col_map.is_empty() {
                    let name_upper = alias_name.to_uppercase();
                    if alias_name != name_upper {
                        dict.insert(name_upper, col_map.clone());
                    }
                    dict.insert(alias_name, col_map);
                }
            }
        }
        ast::TableFactor::NestedJoin {
            table_with_joins, ..
        } => {
            extend_subquery_from_factor(&table_with_joins.relation, flat_schema, dict, dialect);
            for join in &table_with_joins.joins {
                extend_subquery_from_factor(&join.relation, flat_schema, dict, dialect);
            }
        }
        _ => {}
    }
}

/// Expand a table's columns into the CTE dict col_map.
/// If the table is a previously-built CTE, copies from CTE dict.
/// Otherwise expands from schema.
fn expand_table_into_cte_dict(
    table_name: &str,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    col_map: &mut BTreeMap<String, BTreeSet<String>>,
) {
    // Check CTE dict first (handles CTE chains)
    if let Some(prev_cte) = find_cte_in_dict(cte_dict, table_name) {
        for (col, sources) in prev_cte {
            col_map
                .entry(col.clone())
                .or_default()
                .extend(sources.iter().cloned());
        }
        return;
    }

    // Expand from schema
    if let Some(canonical) = find_schema_table(flat_schema, table_name)
        && let Some(cols) = flat_schema.get(canonical)
    {
        for col in cols {
            let col_name = normalize_column_name(col, dialect);
            let source = qualify_ref(canonical, col, dialect);
            col_map.entry(col_name).or_default().insert(source);
        }
    }
}

/// Resolve all column references in an expression to base table sources,
/// using the CTE dict for CTE resolution and schema for base table resolution.
fn resolve_expr_for_cte_dict(
    expr: &Expr,
    alias_map: &AliasMap,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    sources: &mut BTreeSet<String>,
    in_progress: Option<&BTreeMap<String, BTreeSet<String>>>,
) {
    // ValuesOnly now includes CASE/IFF conditions as direct lineage sources.
    // Fallback to All mode only for truly empty expressions (all literals).
    let mut raw_refs = Vec::new();
    extract_columns_from_expr_values_only(expr, &mut raw_refs);

    if raw_refs.is_empty() {
        extract_columns_from_expr(expr, &mut raw_refs);
    }

    for raw in &raw_refs {
        let resolved =
            resolve_ref_with_cte_dict(raw, alias_map, flat_schema, cte_dict, dialect, in_progress);
        sources.extend(resolved);
    }
}

/// Resolve a single column reference using CTE dict + schema (Python's _find_alias_col).
/// `in_progress` allows lateral alias resolution: when building a SELECT's col_map,
/// earlier aliases (e.g., `last_visit_member_a`) are already in the map and can be
/// referenced by later expressions (e.g., `date_diff`) in the same SELECT.
fn resolve_ref_with_cte_dict(
    raw: &RawColumnRef,
    alias_map: &AliasMap,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    in_progress: Option<&BTreeMap<String, BTreeSet<String>>>,
) -> Vec<String> {
    if let Some(table_qualifier) = &raw.table {
        // Qualified reference: resolve table alias first
        let actual_table = alias_map
            .get(table_qualifier.as_str())
            .or_else(|| {
                alias_map
                    .iter()
                    .find(|(k, _)| k.eq_ignore_ascii_case(table_qualifier))
                    .map(|(_, v)| v)
            })
            .cloned()
            .unwrap_or_else(|| table_qualifier.clone());

        // If alias resolves to __subquery__, look up the ORIGINAL qualifier in dict
        let dict_lookup = if actual_table == "__subquery__" {
            table_qualifier.as_str()
        } else {
            actual_table.as_str()
        };

        // Check CTE/subquery dict
        if let Some(cte_cols) = find_cte_in_dict(cte_dict, dict_lookup) {
            let col_key = normalize_column_name(&raw.column, dialect);
            if let Some(cte_sources) = cte_cols.get(&col_key) {
                if !cte_sources.is_empty() {
                    return cte_sources.iter().cloned().collect();
                }
                // Column exists in CTE but has no resolved sources (not in schema).
                // Don't emit the CTE name as a source — return empty to avoid
                // polluting lineage with intermediate CTE/subquery references.
                return vec![];
            }
            // Column not in CTE dict — might be from another table in scope.
            // Fall through to schema check.
        }

        // Check schema
        if actual_table != "__subquery__"
            && let Some(canonical) = find_schema_table(flat_schema, &actual_table)
        {
            return vec![qualify_ref(canonical, &raw.column, dialect)];
        }

        // Fallback: if the source is a CTE or subquery (not a base table),
        // don't emit it as a raw source reference.
        if find_cte_in_dict(cte_dict, dict_lookup).is_some() || actual_table == "__subquery__" {
            return vec![];
        }
        return vec![qualify_ref(dict_lookup, &raw.column, dialect)];
    }

    // Unqualified: search CTE dict first, then schema tables
    let col_key = normalize_column_name(&raw.column, dialect);
    let mut elim_tables = Vec::new();

    for (alias_key, actual_table) in alias_map.iter() {
        if actual_table == "__subquery__" {
            // Try subquery alias in dict
            if let Some(cte_cols) = find_cte_in_dict(cte_dict, alias_key) {
                if let Some(cte_sources) = cte_cols.get(&col_key) {
                    if !cte_sources.is_empty() {
                        return cte_sources.iter().cloned().collect();
                    }
                    // Column exists in subquery output but has no sources
                    // (literal/expression). Return empty — don't fall through
                    // to schema scan which would find a same-named table column.
                    return vec![];
                }
                elim_tables.push(alias_key.clone());
            }
            continue;
        }

        // Check CTE dict
        if let Some(cte_cols) = find_cte_in_dict(cte_dict, actual_table) {
            if let Some(cte_sources) = cte_cols.get(&col_key) {
                if !cte_sources.is_empty() {
                    return cte_sources.iter().cloned().collect();
                }
                // Column in CTE output with no sources (literal). Stop here.
                return vec![];
            }
            elim_tables.push(actual_table.clone());
            continue;
        }

        // Check schema
        if schema_has_column(flat_schema, actual_table, &raw.column) {
            if let Some(canonical) = find_schema_table(flat_schema, actual_table) {
                return vec![qualify_ref(canonical, &raw.column, dialect)];
            }
            return vec![qualify_ref(actual_table, &raw.column, dialect)];
        }
        elim_tables.push(actual_table.clone());
    }

    // Deduction: if all but one table eliminated, use remaining
    let all_tables: Vec<&String> = alias_map
        .values()
        .filter(|t| *t != "__subquery__")
        .collect();
    let remaining: Vec<&&String> = all_tables
        .iter()
        .filter(|t| !elim_tables.contains(t))
        .collect();
    if remaining.len() == 1 {
        return vec![qualify_ref(remaining[0], &raw.column, dialect)];
    }

    // Fallback: search all schema tables
    for table_name in flat_schema.keys() {
        if schema_has_column(flat_schema, table_name, &raw.column) {
            return vec![qualify_ref(table_name, &raw.column, dialect)];
        }
    }

    // Lateral alias resolution: check the in-progress col_map for references
    // to earlier aliases in the same SELECT (Snowflake lateral column alias).
    // e.g., `date_diff` referencing `last_visit_member_a` defined earlier.
    if let Some(ip) = in_progress
        && let Some(lateral_sources) = ip.get(&col_key)
        && !lateral_sources.is_empty()
    {
        return lateral_sources.iter().cloned().collect();
    }

    vec![]
}

// ===========================================================================
// Column reference resolution
// ===========================================================================

/// Build a qualified `"table"."column"` reference string.
fn qualify_ref(table: &str, col: &str, dialect: SqlDialect) -> String {
    let parts: Vec<&str> = table.split('.').collect();
    let mut qualified: Vec<String> = parts.iter().map(|p| format!("\"{}\"", p)).collect();
    qualified.push(format!("\"{}\"", col));
    let result = qualified.join(".");
    match dialect {
        SqlDialect::Snowflake => result.to_uppercase(),
        _ => result,
    }
}

/// Resolve a column reference within a specific scope's alias map.
///
/// For qualified refs: resolve alias only if found in this alias map.
/// For unqualified refs: search FROM tables in this scope only.
///
/// Returns empty if the reference can't be resolved in this scope,
/// allowing the caller to try other scopes.
fn resolve_in_scope(
    raw: &RawColumnRef,
    alias_map: &AliasMap,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> Vec<String> {
    if let Some(table_qualifier) = &raw.table {
        // Look up the qualifier in this scope's alias map
        if let Some(actual_table) = alias_map.get(table_qualifier.as_str()) {
            let actual_table = actual_table.clone();

            if actual_table == "__subquery__" {
                // Subquery alias: resolve through to underlying tables
                return resolve_subquery_column(
                    table_qualifier,
                    &raw.column,
                    subquery_sources,
                    flat_schema,
                    dialect,
                );
            }

            if cte_names.contains(&actual_table) {
                return resolve_cte_column(
                    &actual_table,
                    &raw.column,
                    cte_source_tables,
                    flat_schema,
                    dialect,
                );
            }

            return vec![qualify_ref(&actual_table, &raw.column, dialect)];
        }

        // Qualifier not found in this scope — return empty so caller tries next scope
        return vec![];
    }

    // Unqualified column: search FROM tables in this scope
    let from_tables: BTreeSet<&String> = alias_map
        .values()
        .filter(|t| !cte_names.contains(*t) && *t != "__subquery__")
        .collect();

    for table_name in &from_tables {
        if schema_has_column(flat_schema, table_name, &raw.column) {
            return vec![qualify_ref(table_name, &raw.column, dialect)];
        }
    }

    // Search all schema tables as fallback
    for table_name in flat_schema.keys() {
        if schema_has_column(flat_schema, table_name, &raw.column) {
            return vec![qualify_ref(table_name, &raw.column, dialect)];
        }
    }

    vec![]
}

/// Resolve a column reference with fallback (for single-scope contexts like
/// supplement_missing_columns where we know the correct alias map).
fn resolve_column_ref(
    raw: &RawColumnRef,
    alias_map: &AliasMap,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> Vec<String> {
    if let Some(table_qualifier) = &raw.table {
        let actual_table = alias_map
            .get(table_qualifier.as_str())
            .cloned()
            .unwrap_or_else(|| table_qualifier.clone());

        if actual_table == "__subquery__" {
            return resolve_subquery_column(
                table_qualifier,
                &raw.column,
                subquery_sources,
                flat_schema,
                dialect,
            );
        }

        if cte_names.contains(&actual_table) {
            return resolve_cte_column(
                &actual_table,
                &raw.column,
                cte_source_tables,
                flat_schema,
                dialect,
            );
        }

        return vec![qualify_ref(&actual_table, &raw.column, dialect)];
    }

    // Unqualified: search FROM tables then all schema tables
    let from_tables: BTreeSet<&String> = alias_map
        .values()
        .filter(|t| !cte_names.contains(*t) && *t != "__subquery__")
        .collect();

    for table_name in &from_tables {
        if schema_has_column(flat_schema, table_name, &raw.column) {
            return vec![qualify_ref(table_name, &raw.column, dialect)];
        }
    }

    for table_name in flat_schema.keys() {
        if schema_has_column(flat_schema, table_name, &raw.column) {
            return vec![qualify_ref(table_name, &raw.column, dialect)];
        }
    }

    vec![]
}

/// Resolve a CTE column reference to base table column(s).
fn resolve_cte_column(
    cte_name: &str,
    column: &str,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) -> Vec<String> {
    resolve_cte_column_inner(cte_name, column, cte_source_tables, flat_schema, dialect, 0)
}

fn resolve_cte_column_inner(
    cte_name: &str,
    column: &str,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<String> {
    if depth > 10 {
        return vec![];
    }

    // Case-insensitive lookup: Snowflake uppercases CTE names in star sources
    let source_tables = cte_source_tables.get(cte_name).or_else(|| {
        let lower = cte_name.to_lowercase();
        cte_source_tables
            .iter()
            .find(|(k, _)| k.to_lowercase() == lower)
            .map(|(_, v)| v)
    });

    if let Some(source_tables) = source_tables {
        // First pass: check schema tables directly
        for table in source_tables {
            if schema_has_column(flat_schema, table, column) {
                if let Some(canonical) = find_schema_table(flat_schema, table) {
                    return vec![qualify_ref(canonical, column, dialect)];
                }
                return vec![qualify_ref(table, column, dialect)];
            }
        }
        // Second pass: recurse through CTE sources (CTE → CTE chains)
        for table in source_tables {
            let is_cte = cte_source_tables.contains_key(table.as_str())
                || cte_source_tables
                    .keys()
                    .any(|k| k.eq_ignore_ascii_case(table));
            if is_cte {
                let resolved = resolve_cte_column_inner(
                    table,
                    column,
                    cte_source_tables,
                    flat_schema,
                    dialect,
                    depth + 1,
                );
                if !resolved.is_empty() {
                    return resolved;
                }
            }
        }
        // Single source: assume passthrough ONLY if the source table is in
        // the schema (handles incomplete schemas where column isn't explicitly
        // listed but still belongs to the table).
        if source_tables.len() == 1
            && let Some(canonical) = find_schema_table(flat_schema, &source_tables[0])
        {
            return vec![qualify_ref(canonical, column, dialect)];
        }
    }
    vec![]
}

/// Resolve a subquery alias column to its underlying table.
/// Recursively resolves through nested subquery aliases.
fn resolve_subquery_column(
    subquery_alias: &str,
    column: &str,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) -> Vec<String> {
    resolve_subquery_column_inner(
        subquery_alias,
        column,
        subquery_sources,
        flat_schema,
        dialect,
        0,
    )
}

fn resolve_subquery_column_inner(
    subquery_alias: &str,
    column: &str,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<String> {
    if depth > 5 {
        return vec![];
    }
    let inner_tables = subquery_sources.get(subquery_alias).or_else(|| {
        subquery_sources
            .iter()
            .find(|(k, _)| k.eq_ignore_ascii_case(subquery_alias))
            .map(|(_, v)| v)
    });

    if let Some(inner_tables) = inner_tables {
        for table in inner_tables {
            // Direct schema table match
            if schema_has_column(flat_schema, table, column) {
                if let Some(canonical) = find_schema_table(flat_schema, table) {
                    return vec![qualify_ref(canonical, column, dialect)];
                }
                return vec![qualify_ref(table, column, dialect)];
            }
            // Recursive: inner table might be another subquery alias or CTE
            if subquery_sources.contains_key(table.as_str())
                || subquery_sources
                    .keys()
                    .any(|k| k.eq_ignore_ascii_case(table))
            {
                let r = resolve_subquery_column_inner(
                    table,
                    column,
                    subquery_sources,
                    flat_schema,
                    dialect,
                    depth + 1,
                );
                if !r.is_empty() {
                    return r;
                }
            }
        }
    }
    vec![]
}

// ===========================================================================
// Column reference extraction from AST expressions
// ===========================================================================

/// Controls which column references are extracted from expressions.
#[derive(Clone, Copy, PartialEq, Eq)]
enum ExtractionMode {
    /// Extract ALL column references (conditions, values, window clauses).
    All,
    /// Extract only value-producing references:
    /// - CASE: skip operand and conditions, keep results and else
    /// - Window functions: skip PARTITION BY and ORDER BY, keep function args
    ValuesOnly,
}

fn extract_columns_from_expr(expr: &Expr, refs: &mut Vec<RawColumnRef>) {
    extract_columns_impl(expr, refs, ExtractionMode::All);
}

/// Extract column refs that appear in PARTITION BY clauses of window functions.
/// Used to preserve PARTITION BY dependencies in complete mode while still
/// stripping CASE condition columns.
fn extract_window_partition_refs(expr: &Expr, refs: &mut Vec<RawColumnRef>) {
    match expr {
        Expr::Function(f) => {
            // Extract PARTITION BY refs from this window function
            if let Some(ast::WindowType::WindowSpec(spec)) = &f.over {
                for partition_expr in &spec.partition_by {
                    extract_columns_impl(partition_expr, refs, ExtractionMode::All);
                }
            }
            // Also recurse into function args for nested window functions
            if let ast::FunctionArguments::List(arg_list) = &f.args {
                for arg in &arg_list.args {
                    match arg {
                        ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e))
                        | ast::FunctionArg::Named {
                            arg: ast::FunctionArgExpr::Expr(e),
                            ..
                        } => extract_window_partition_refs(e, refs),
                        _ => {}
                    }
                }
            }
        }
        Expr::Nested(inner) => extract_window_partition_refs(inner, refs),
        Expr::Case {
            operand,
            conditions,
            results,
            else_result,
            ..
        } => {
            if let Some(op) = operand {
                extract_window_partition_refs(op, refs);
            }
            for cond in conditions {
                extract_window_partition_refs(cond, refs);
            }
            for result in results {
                extract_window_partition_refs(result, refs);
            }
            if let Some(el) = else_result {
                extract_window_partition_refs(el, refs);
            }
        }
        Expr::Cast { expr: e, .. } => {
            extract_window_partition_refs(e, refs);
        }
        _ => {}
    }
}

fn extract_columns_from_expr_values_only(expr: &Expr, refs: &mut Vec<RawColumnRef>) {
    extract_columns_impl(expr, refs, ExtractionMode::ValuesOnly);
}

/// Check if an expression is a bare identifier matching a known date part keyword.
/// These appear as first args in DATE_PART, DATEDIFF, etc. and should not be
/// treated as column references (e.g., DATE_PART(EPOCH_MILLISECOND, col)).
fn is_date_part_keyword(expr: &Expr) -> bool {
    let ident = match expr {
        Expr::Identifier(id) => id.value.to_uppercase(),
        // String literals like 'day' are always safe to skip
        Expr::Value(ast::Value::SingleQuotedString(_)) => return true,
        _ => return false,
    };
    // Comprehensive list of Snowflake date/time part keywords
    matches!(
        ident.as_str(),
        "YEAR"
            | "YEARS"
            | "Y"
            | "YY"
            | "YYY"
            | "YYYY"
            | "MONTH"
            | "MONTHS"
            | "MM"
            | "MON"
            | "DAY"
            | "DAYS"
            | "D"
            | "DD"
            | "DAYOFMONTH"
            | "DAYOFWEEK"
            | "DOW"
            | "DW"
            | "DAYOFWEEKISO"
            | "DOWISO"
            | "DW_ISO"
            | "DAYOFYEAR"
            | "DOY"
            | "DY"
            | "WEEK"
            | "WEEKS"
            | "W"
            | "WK"
            | "WEEKISO"
            | "WEEKOFYEAR"
            | "WOY"
            | "WY"
            | "QUARTER"
            | "QUARTERS"
            | "Q"
            | "QTR"
            | "HOUR"
            | "HOURS"
            | "H"
            | "HH"
            | "MINUTE"
            | "MINUTES"
            | "M"
            | "MI"
            | "SECOND"
            | "SECONDS"
            | "S"
            | "SS"
            | "SEC"
            | "MILLISECOND"
            | "MILLISECONDS"
            | "MS"
            | "MSEC"
            | "MICROSECOND"
            | "MICROSECONDS"
            | "US"
            | "USEC"
            | "NANOSECOND"
            | "NANOSECONDS"
            | "NS"
            | "NSEC"
            | "EPOCH"
            | "EPOCH_SECOND"
            | "EPOCH_MILLISECOND"
            | "EPOCH_MICROSECOND"
            | "EPOCH_NANOSECOND"
            | "TIMEZONE"
            | "TIMEZONE_HOUR"
            | "TIMEZONE_MINUTE"
            | "TZH"
            | "TZM"
    )
}

fn extract_columns_impl(expr: &Expr, refs: &mut Vec<RawColumnRef>, mode: ExtractionMode) {
    match expr {
        Expr::Identifier(ident) => {
            refs.push(RawColumnRef {
                table: None,
                column: ident.value.clone(),
            });
        }
        Expr::CompoundIdentifier(parts) => {
            if parts.len() == 2 {
                refs.push(RawColumnRef {
                    table: Some(parts[0].value.clone()),
                    column: parts[1].value.clone(),
                });
            } else if parts.len() >= 3 {
                let col = parts.last().unwrap().value.clone();
                let table_parts: Vec<_> = parts[..parts.len() - 1]
                    .iter()
                    .map(|p| p.value.clone())
                    .collect();
                refs.push(RawColumnRef {
                    table: Some(table_parts.join(".")),
                    column: col,
                });
            }
        }
        Expr::BinaryOp { left, right, .. } => {
            extract_columns_impl(left, refs, mode);
            extract_columns_impl(right, refs, mode);
        }
        Expr::UnaryOp { expr: inner, .. } => {
            extract_columns_impl(inner, refs, mode);
        }
        Expr::IsNull(inner)
        | Expr::IsNotNull(inner)
        | Expr::IsTrue(inner)
        | Expr::IsFalse(inner)
        | Expr::IsNotTrue(inner)
        | Expr::IsNotFalse(inner)
        | Expr::IsUnknown(inner)
        | Expr::IsNotUnknown(inner) => {
            extract_columns_impl(inner, refs, mode);
        }
        Expr::Between {
            expr: e, low, high, ..
        } => {
            extract_columns_impl(e, refs, mode);
            extract_columns_impl(low, refs, mode);
            extract_columns_impl(high, refs, mode);
        }
        Expr::InList { expr: e, list, .. } => {
            extract_columns_impl(e, refs, mode);
            for item in list {
                extract_columns_impl(item, refs, mode);
            }
        }
        Expr::InSubquery { expr: e, .. } => {
            extract_columns_impl(e, refs, mode);
            // Don't recurse into subquery here — handled by scope-paired extraction
        }
        Expr::Exists { .. } => {
            // Don't recurse — handled by scope-paired extraction
        }
        Expr::Subquery(_) => {
            // Don't recurse — handled by scope-paired extraction
        }
        Expr::Nested(inner) => {
            extract_columns_impl(inner, refs, mode);
        }
        Expr::Function(f) => {
            let func_name = f
                .name
                .0
                .last()
                .map(|i| i.value.to_uppercase())
                .unwrap_or_default();

            // IFF(cond, then, else) and DECODE — Snowflake conditional functions.
            // The condition arg is a direct lineage source: it determines which
            // result branch is selected, so include it in all extraction modes.
            // Date/time functions where the first arg may be a date part keyword
            // (e.g., DATE_PART(EPOCH_MILLISECOND, col), DATEDIFF(DAY, a, b)).
            // Only skip if the first arg is a bare identifier matching a known
            // date part keyword — never skip complex expressions or qualified refs.
            static DATE_PART_FUNCS: &[&str] = &[
                "DATE_PART",
                "DATEPART",
                "EXTRACT",
                "DATE_TRUNC",
                "DATETRUNC",
                "DATEDIFF",
                "DATEADD",
                "DATE_ADD",
                "DATE_SUB",
                "TIMESTAMPADD",
                "TIMESTAMPDIFF",
                "TIME_SLICE",
            ];
            let is_date_part_func = DATE_PART_FUNCS.contains(&func_name.as_str());

            match &f.args {
                ast::FunctionArguments::List(arg_list) => {
                    for (idx, arg) in arg_list.args.iter().enumerate() {
                        // Skip date part keyword: only if first arg is a bare identifier
                        // matching a known date part name (YEAR, MONTH, DAY, EPOCH_*, etc.)
                        if idx == 0
                            && is_date_part_func
                            && let ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e)) = arg
                            && is_date_part_keyword(e)
                        {
                            continue;
                        }
                        match arg {
                            ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e)) => {
                                extract_columns_impl(e, refs, mode);
                            }
                            ast::FunctionArg::Named {
                                arg: ast::FunctionArgExpr::Expr(e),
                                ..
                            } => {
                                extract_columns_impl(e, refs, mode);
                            }
                            _ => {}
                        }
                    }
                }
                ast::FunctionArguments::Subquery(_) => {
                    // Subquery in function args — don't recurse
                }
                ast::FunctionArguments::None => {}
            }
            // WITHIN GROUP (ORDER BY ...) — aggregate ordering clauses.
            // Snowflake ACCESS_USAGE treats these as direct lineage sources
            // (they influence the output value, e.g. LISTAGG ordering,
            // ARRAY_AGG ordering). Extract in all modes.
            for ob in &f.within_group {
                extract_columns_impl(&ob.expr, refs, mode);
            }

            // Inline ORDER BY inside function argument list
            // (e.g., ARRAY_AGG(x ORDER BY y) — BigQuery/Snowflake syntax).
            if let ast::FunctionArguments::List(arg_list) = &f.args {
                for clause in &arg_list.clauses {
                    if let ast::FunctionArgumentClause::OrderBy(order_exprs) = clause {
                        for ob in order_exprs {
                            extract_columns_impl(&ob.expr, refs, mode);
                        }
                    }
                }
            }

            // Window PARTITION BY and ORDER BY handling by extraction mode:
            // - All: include both PARTITION BY and ORDER BY
            // - ValuesOnly: skip both
            //   Neither PARTITION BY nor ORDER BY produce data in the output —
            //   they control grouping and ordering but don't contribute values.
            //   Snowflake ACCESS_USAGE excludes both from column lineage.
            if let Some(ast::WindowType::WindowSpec(spec)) = &f.over
                && mode == ExtractionMode::All
            {
                for expr in &spec.partition_by {
                    extract_columns_impl(expr, refs, mode);
                }
                for ob in &spec.order_by {
                    extract_columns_impl(&ob.expr, refs, mode);
                }
            }
        }
        Expr::Cast { expr: e, .. } => {
            extract_columns_impl(e, refs, mode);
        }
        Expr::Case {
            operand,
            conditions,
            results,
            else_result,
            ..
        } => {
            // CASE conditions are direct lineage sources: they determine which
            // result branch is selected. Include in all extraction modes.
            if let Some(op) = operand {
                extract_columns_impl(op, refs, mode);
            }
            for cond in conditions {
                extract_columns_impl(cond, refs, mode);
            }
            for res in results {
                extract_columns_impl(res, refs, mode);
            }
            if let Some(el) = else_result {
                extract_columns_impl(el, refs, mode);
            }
        }
        Expr::Like { expr, pattern, .. } | Expr::ILike { expr, pattern, .. } => {
            extract_columns_impl(expr, refs, mode);
            extract_columns_impl(pattern, refs, mode);
        }
        Expr::JsonAccess { value, .. } => {
            // Snowflake variant access: col:field, col[0]:field, col:a:b
            // The base column is in `value`; the path selects sub-fields.
            // For lineage, the data source is the base column (the VARIANT column).
            extract_columns_impl(value, refs, mode);
        }
        // Special syntax expressions — sqlparser represents these as dedicated
        // enum variants rather than Expr::Function. Without these arms, the
        // catch-all `_ => {}` silently drops all column references inside them.
        Expr::Substring {
            expr,
            substring_from,
            substring_for,
            ..
        } => {
            extract_columns_impl(expr, refs, mode);
            if let Some(from_expr) = substring_from {
                extract_columns_impl(from_expr, refs, mode);
            }
            if let Some(for_expr) = substring_for {
                extract_columns_impl(for_expr, refs, mode);
            }
        }
        Expr::Trim {
            expr,
            trim_what,
            trim_characters,
            ..
        } => {
            extract_columns_impl(expr, refs, mode);
            if let Some(what) = trim_what {
                extract_columns_impl(what, refs, mode);
            }
            if let Some(chars) = trim_characters {
                for c in chars {
                    extract_columns_impl(c, refs, mode);
                }
            }
        }
        Expr::Overlay {
            expr,
            overlay_what,
            overlay_from,
            overlay_for,
            ..
        } => {
            extract_columns_impl(expr, refs, mode);
            extract_columns_impl(overlay_what, refs, mode);
            extract_columns_impl(overlay_from, refs, mode);
            if let Some(for_expr) = overlay_for {
                extract_columns_impl(for_expr, refs, mode);
            }
        }
        Expr::Position { expr, r#in } => {
            extract_columns_impl(expr, refs, mode);
            extract_columns_impl(r#in, refs, mode);
        }
        Expr::Extract { expr, .. } => {
            extract_columns_impl(expr, refs, mode);
        }
        Expr::Ceil { expr, .. } | Expr::Floor { expr, .. } => {
            extract_columns_impl(expr, refs, mode);
        }
        Expr::Collate { expr, .. } => {
            extract_columns_impl(expr, refs, mode);
        }
        Expr::IsDistinctFrom(a, b) | Expr::IsNotDistinctFrom(a, b) => {
            extract_columns_impl(a, refs, mode);
            extract_columns_impl(b, refs, mode);
        }
        Expr::AtTimeZone {
            timestamp,
            time_zone,
        } => {
            extract_columns_impl(timestamp, refs, mode);
            extract_columns_impl(time_zone, refs, mode);
        }
        Expr::SimilarTo { expr, pattern, .. } | Expr::RLike { expr, pattern, .. } => {
            extract_columns_impl(expr, refs, mode);
            extract_columns_impl(pattern, refs, mode);
        }
        Expr::InUnnest {
            expr, array_expr, ..
        } => {
            extract_columns_impl(expr, refs, mode);
            extract_columns_impl(array_expr, refs, mode);
        }
        Expr::AnyOp { left, right, .. } | Expr::AllOp { left, right, .. } => {
            extract_columns_impl(left, refs, mode);
            extract_columns_impl(right, refs, mode);
        }
        Expr::Tuple(exprs) => {
            for e in exprs {
                extract_columns_impl(e, refs, mode);
            }
        }
        Expr::Prior(inner) => {
            extract_columns_impl(inner, refs, mode);
        }
        Expr::Named { expr, .. } => {
            extract_columns_impl(expr, refs, mode);
        }
        Expr::CompoundFieldAccess { root, .. } => {
            extract_columns_impl(root, refs, mode);
        }
        Expr::IsNormalized { expr, .. } => {
            extract_columns_impl(expr, refs, mode);
        }
        _ => {}
    }
}

fn is_numeric_literal(expr: &Expr) -> bool {
    matches!(expr, Expr::Value(ast::Value::Number(_, _)))
}

// ===========================================================================
// Scope-paired shared condition extraction and resolution
//
// Each SQL scope (CTE body, subquery, main query) has its own alias map.
// We extract raw column refs from each scope's shared conditions and resolve
// them using THAT scope's alias map — never mixing scopes.
// ===========================================================================

/// Context passed through the scope-paired extraction.
struct ResolutionCtx<'a> {
    flat_schema: &'a FlatSchema,
    cte_names: &'a BTreeSet<String>,
    cte_source_tables: &'a BTreeMap<String, Vec<String>>,
    subquery_sources: &'a BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
}

/// Extract and resolve shared conditions from the entire query, scope by scope.
fn extract_and_resolve_shared_deps(
    query: &ast::Query,
    used_ctes: &BTreeSet<String>,
    _subquery_sources: &BTreeMap<String, Vec<String>>,
    cte_alias_maps: &BTreeMap<String, AliasMap>,
    ctx: &ResolutionCtx,
) -> BTreeSet<String> {
    let mut deps = BTreeSet::new();

    // Process used CTE bodies with their own alias maps
    if let Some(with) = &query.with {
        for cte in &with.cte_tables {
            let name = cte.alias.name.value.clone();
            if !used_ctes.contains(&name) {
                continue;
            }
            let alias_map = cte_alias_maps.get(&name).cloned().unwrap_or_default();
            resolve_shared_from_set_expr(&cte.query.body, &alias_map, ctx, &mut deps);
        }
    }

    // Process main query body with its alias map
    let mut sub = BTreeMap::new();
    let alias_map = build_alias_map_from_body(&query.body, &mut sub);
    resolve_shared_from_set_expr_with_query(&query.body, &alias_map, ctx, Some(query), &mut deps);

    deps
}

fn resolve_shared_from_set_expr(
    body: &SetExpr,
    alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    deps: &mut BTreeSet<String>,
) {
    resolve_shared_from_set_expr_with_query(body, alias_map, ctx, None, deps);
}

fn resolve_shared_from_set_expr_with_query(
    body: &SetExpr,
    alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    query: Option<&ast::Query>,
    deps: &mut BTreeSet<String>,
) {
    match body {
        SetExpr::Select(sel) => {
            resolve_shared_from_select(sel, alias_map, ctx, query, deps);
        }
        SetExpr::SetOperation { left, right, .. } => {
            resolve_shared_from_set_expr_with_query(left, alias_map, ctx, query, deps);
            let mut sub = BTreeMap::new();
            let right_alias = build_alias_map_from_body(right, &mut sub);
            resolve_shared_from_set_expr_with_query(right, &right_alias, ctx, query, deps);
        }
        SetExpr::Query(q) => {
            let mut sub = BTreeMap::new();
            let inner_alias = build_alias_map_from_body(&q.body, &mut sub);
            resolve_shared_from_set_expr_with_query(&q.body, &inner_alias, ctx, query, deps);
        }
        _ => {}
    }
}

/// Resolve raw refs using a specific scope's alias map.
/// `query` is optionally provided for CTE expression column fallback.
fn resolve_refs_in_scope(
    refs: &[RawColumnRef],
    alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    deps: &mut BTreeSet<String>,
) {
    resolve_refs_in_scope_with_query(refs, alias_map, ctx, None, deps);
}

fn resolve_refs_in_scope_with_query(
    refs: &[RawColumnRef],
    alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    query: Option<&ast::Query>,
    deps: &mut BTreeSet<String>,
) {
    for raw in refs {
        let resolved = resolve_in_scope(
            raw,
            alias_map,
            ctx.flat_schema,
            ctx.cte_names,
            ctx.cte_source_tables,
            ctx.subquery_sources,
            ctx.dialect,
        );
        if !resolved.is_empty() {
            deps.extend(resolved);
        } else if let Some(query) = query {
            // Fallback: check if this is a CTE expression column (e.g., `rank` from RANK())
            // For unqualified refs: check all CTEs referenced in alias_map
            // For qualified refs: check the specific CTE
            let cte_to_check: Vec<String> = if let Some(table) = &raw.table {
                if ctx.cte_names.contains(table) {
                    vec![table.clone()]
                } else if let Some(actual) = alias_map.get(table.as_str()) {
                    if ctx.cte_names.contains(actual) {
                        vec![actual.clone()]
                    } else {
                        vec![]
                    }
                } else {
                    vec![]
                }
            } else {
                // Unqualified: try all CTEs referenced by tables in alias_map
                alias_map
                    .values()
                    .filter(|v| ctx.cte_names.contains(*v))
                    .cloned()
                    .collect()
            };

            for cte_name in &cte_to_check {
                let expr_deps = resolve_cte_expression_column(
                    cte_name,
                    &raw.column,
                    query,
                    ctx.flat_schema,
                    ctx.cte_names,
                    ctx.cte_source_tables,
                    ctx.subquery_sources,
                    ctx.dialect,
                );
                if !expr_deps.is_empty() {
                    deps.extend(expr_deps);
                    break;
                }
            }
        }
    }
}

fn resolve_shared_from_select(
    sel: &ast::Select,
    alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    query: Option<&ast::Query>,
    deps: &mut BTreeSet<String>,
) {
    // WHERE — use query-aware resolution for CTE expression column fallback
    if let Some(where_expr) = &sel.selection {
        let mut refs = Vec::new();
        extract_columns_from_expr(where_expr, &mut refs);
        resolve_refs_in_scope_with_query(&refs, alias_map, ctx, query, deps);

        // Recurse into subqueries in WHERE with their own alias maps
        // Pass outer alias_map so correlated refs (t1.a in EXISTS WHERE t1.a = t2.a) resolve
        extract_subquery_deps_from_expr(where_expr, alias_map, ctx, deps);
    }

    // HAVING
    if let Some(having_expr) = &sel.having {
        let mut refs = Vec::new();
        extract_columns_from_expr(having_expr, &mut refs);
        resolve_refs_in_scope_with_query(&refs, alias_map, ctx, query, deps);
    }

    // GROUP BY (non-positional)
    match &sel.group_by {
        GroupByExpr::Expressions(exprs, _) => {
            for expr in exprs {
                if !is_numeric_literal(expr) {
                    let mut refs = Vec::new();
                    extract_columns_from_expr(expr, &mut refs);
                    resolve_refs_in_scope(&refs, alias_map, ctx, deps);
                }
            }
        }
        GroupByExpr::All(_) => {}
    }

    // JOIN ON / USING
    for twj in &sel.from {
        for join in &twj.joins {
            match &join.join_operator {
                JoinOperator::Inner(constraint)
                | JoinOperator::LeftOuter(constraint)
                | JoinOperator::RightOuter(constraint)
                | JoinOperator::FullOuter(constraint) => match constraint {
                    JoinConstraint::On(on_expr) => {
                        let mut refs = Vec::new();
                        extract_columns_from_expr(on_expr, &mut refs);
                        resolve_refs_in_scope(&refs, alias_map, ctx, deps);
                    }
                    JoinConstraint::Using(cols) => {
                        // USING(a) creates a shared condition on 'a' for ALL joined tables
                        for col in cols {
                            let col_name =
                                col.0.last().map(|i| i.value.clone()).unwrap_or_default();

                            // Resolve against every table in the FROM scope that has this column
                            let mut resolved_any = false;
                            for (alias, actual) in alias_map.iter() {
                                if actual == "__subquery__" {
                                    continue;
                                }
                                // Check if this table has the column
                                let table_to_check = if ctx.cte_names.contains(actual) {
                                    // CTE: check its source tables
                                    if let Some(sources) = ctx.cte_source_tables.get(actual) {
                                        sources.iter().any(|s| {
                                            schema_has_column(ctx.flat_schema, s, &col_name)
                                        })
                                    } else {
                                        false
                                    }
                                } else {
                                    schema_has_column(ctx.flat_schema, actual, &col_name)
                                };

                                if table_to_check {
                                    let raw = RawColumnRef {
                                        table: Some(alias.clone()),
                                        column: col_name.clone(),
                                    };
                                    let resolved = resolve_in_scope(
                                        &raw,
                                        alias_map,
                                        ctx.flat_schema,
                                        ctx.cte_names,
                                        ctx.cte_source_tables,
                                        ctx.subquery_sources,
                                        ctx.dialect,
                                    );
                                    deps.extend(resolved);
                                    resolved_any = true;
                                }
                            }
                            // Fallback: unqualified resolution
                            if !resolved_any {
                                let raw = RawColumnRef {
                                    table: None,
                                    column: col_name,
                                };
                                let resolved = resolve_in_scope(
                                    &raw,
                                    alias_map,
                                    ctx.flat_schema,
                                    ctx.cte_names,
                                    ctx.cte_source_tables,
                                    ctx.subquery_sources,
                                    ctx.dialect,
                                );
                                deps.extend(resolved);
                            }
                        }
                    }
                    _ => {}
                },
                _ => {}
            }
        }

        // Subqueries in FROM: resolve their shared conditions in their own scope.
        // Pass outer alias_map so LATERAL subqueries can resolve correlated refs.
        resolve_subquery_factor_deps(&twj.relation, alias_map, ctx, deps);
        for join in &twj.joins {
            resolve_subquery_factor_deps(&join.relation, alias_map, ctx, deps);
        }
    }

    // Scalar subqueries in SELECT list: their WHERE conditions are shared
    for item in &sel.projection {
        match item {
            SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
                extract_subquery_deps_from_expr(e, alias_map, ctx, deps);
            }
            _ => {}
        }
    }
}

/// Recursively find subqueries in an expression and extract their shared conditions.
/// `outer_alias_map` is the enclosing scope's alias map for resolving correlated refs.
fn extract_subquery_deps_from_expr(
    expr: &Expr,
    outer_alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    deps: &mut BTreeSet<String>,
) {
    match expr {
        Expr::InSubquery { subquery, .. } => {
            // IN subquery: extract both SELECT refs and WHERE conditions
            let mut sub = BTreeMap::new();
            let subq_alias = build_alias_map_from_body(&subquery.body, &mut sub);

            // Build a merged alias map: subquery's own + outer scope for correlated refs
            let mut merged = outer_alias_map.clone();
            merged.extend(subq_alias.clone());

            // SELECT list columns from the subquery
            if let SetExpr::Select(sel) = subquery.body.as_ref() {
                for item in &sel.projection {
                    match item {
                        SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
                            let mut refs = Vec::new();
                            extract_columns_from_expr(e, &mut refs);
                            resolve_refs_in_scope(&refs, &merged, ctx, deps);
                        }
                        _ => {}
                    }
                }
            }

            // Shared conditions within the subquery (using merged scope)
            resolve_shared_from_set_expr(&subquery.body, &merged, ctx, deps);
        }
        Expr::Exists { subquery, .. } => {
            // EXISTS: correlation conditions resolve against both inner and outer scope
            let mut sub = BTreeMap::new();
            let subq_alias = build_alias_map_from_body(&subquery.body, &mut sub);
            let mut merged = outer_alias_map.clone();
            merged.extend(subq_alias);
            resolve_shared_from_set_expr(&subquery.body, &merged, ctx, deps);
        }
        Expr::Subquery(subquery) => {
            // Scalar subquery: its WHERE conditions are shared
            let mut sub = BTreeMap::new();
            let subq_alias = build_alias_map_from_body(&subquery.body, &mut sub);
            let mut merged = outer_alias_map.clone();
            merged.extend(subq_alias);
            resolve_shared_from_set_expr(&subquery.body, &merged, ctx, deps);
        }
        // Recurse into compound expressions looking for subqueries
        Expr::BinaryOp { left, right, .. } => {
            extract_subquery_deps_from_expr(left, outer_alias_map, ctx, deps);
            extract_subquery_deps_from_expr(right, outer_alias_map, ctx, deps);
        }
        Expr::UnaryOp { expr: inner, .. } | Expr::Nested(inner) => {
            extract_subquery_deps_from_expr(inner, outer_alias_map, ctx, deps);
        }
        Expr::Function(f) => {
            if let ast::FunctionArguments::List(arg_list) = &f.args {
                for arg in &arg_list.args {
                    match arg {
                        ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e))
                        | ast::FunctionArg::Named {
                            arg: ast::FunctionArgExpr::Expr(e),
                            ..
                        } => {
                            extract_subquery_deps_from_expr(e, outer_alias_map, ctx, deps);
                        }
                        _ => {}
                    }
                }
            }
        }
        _ => {}
    }
}

/// Resolve shared conditions from a subquery table factor.
fn resolve_subquery_factor_deps(
    factor: &ast::TableFactor,
    outer_alias_map: &AliasMap,
    ctx: &ResolutionCtx,
    deps: &mut BTreeSet<String>,
) {
    match factor {
        ast::TableFactor::Derived {
            subquery, lateral, ..
        } => {
            let mut sub = BTreeMap::new();
            let subq_alias = build_alias_map_from_body(&subquery.body, &mut sub);
            if *lateral {
                // LATERAL subqueries can reference outer tables — merge alias maps
                // so correlated refs (e.g., `table1.a` in `WHERE table1.a = table2.a`) resolve.
                let mut merged = outer_alias_map.clone();
                merged.extend(subq_alias);
                resolve_shared_from_set_expr(&subquery.body, &merged, ctx, deps);
            } else {
                resolve_shared_from_set_expr(&subquery.body, &subq_alias, ctx, deps);
            }
        }
        ast::TableFactor::Function { args, .. } => {
            // LATERAL FLATTEN and similar table-valued functions: their arguments
            // reference columns that become shared dependencies (e.g., `flatten(input => t.c)`
            // makes `t.c` a dependency of all output columns).
            let mut refs = Vec::new();
            for arg in args {
                match arg {
                    ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e))
                    | ast::FunctionArg::Named {
                        arg: ast::FunctionArgExpr::Expr(e),
                        ..
                    }
                    | ast::FunctionArg::ExprNamed {
                        arg: ast::FunctionArgExpr::Expr(e),
                        ..
                    } => {
                        extract_columns_from_expr(e, &mut refs);
                    }
                    _ => {}
                }
            }
            if !refs.is_empty() {
                resolve_refs_in_scope(&refs, outer_alias_map, ctx, deps);
            }
        }
        _ => {}
    }
}

// ===========================================================================
// GROUP BY positional reference resolution
// ===========================================================================

fn resolve_group_by_positional(
    query: &ast::Query,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> BTreeSet<String> {
    let mut deps = BTreeSet::new();

    let (select_items, alias_map) = match &*query.body {
        SetExpr::Select(sel) => {
            let mut sub = BTreeMap::new();
            (sel.projection.clone(), build_alias_map(&sel.from, &mut sub))
        }
        _ => return deps,
    };

    let group_by_exprs = match query.body.as_ref() {
        SetExpr::Select(sel) => match &sel.group_by {
            GroupByExpr::Expressions(exprs, _) => exprs.clone(),
            _ => return deps,
        },
        _ => return deps,
    };

    for gb_expr in &group_by_exprs {
        if let Expr::Value(ast::Value::Number(n, _)) = gb_expr
            && let Ok(pos) = n.parse::<usize>()
            && pos >= 1
            && pos <= select_items.len()
        {
            let item = &select_items[pos - 1];
            let mut item_refs = Vec::new();
            match item {
                SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
                    extract_columns_from_expr(e, &mut item_refs);
                }
                _ => {}
            }
            for raw in &item_refs {
                let resolved = resolve_column_ref(
                    raw,
                    &alias_map,
                    flat_schema,
                    cte_names,
                    cte_source_tables,
                    subquery_sources,
                    dialect,
                );
                deps.extend(resolved);
            }
        }
    }

    deps
}

// ===========================================================================
// SELECT list analysis (for aggregates and expressions)
// ===========================================================================

fn select_item_name(item: &SelectItem, position: usize, dialect: SqlDialect) -> Option<String> {
    match item {
        SelectItem::ExprWithAlias { alias, .. } => {
            Some(normalize_column_name(&alias.value, dialect))
        }
        SelectItem::UnnamedExpr(expr) => match expr {
            Expr::Identifier(ident) => Some(normalize_column_name(&ident.value, dialect)),
            Expr::CompoundIdentifier(parts) => parts
                .last()
                .map(|p| normalize_column_name(&p.value, dialect)),
            _ => Some(format!("_col_{}", position)),
        },
        SelectItem::Wildcard(_) => Some(normalize_column_name("*", dialect)),
        SelectItem::QualifiedWildcard(name, _) => Some(format!("{}.*", name)),
    }
}

/// Add or merge columns from AST extraction that polyglot didn't track.
/// For columns already in column_dict (from polyglot edges), merges additional
/// sources (e.g., CASE condition columns). For new columns, creates entries.
fn supplement_missing_columns(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    for (i, item) in info.select_items.iter().enumerate() {
        let col_name = match select_item_name(item, i, dialect) {
            Some(n) => n,
            None => continue,
        };

        let expr = match item {
            SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => e,
            _ => continue,
        };

        let mut sources = BTreeSet::new();

        // For scalar subqueries, resolve using the subquery's own alias map
        if let Expr::Subquery(subquery) = expr {
            resolve_scalar_subquery_columns(
                subquery,
                flat_schema,
                &info.cte_names,
                &info.cte_source_tables,
                &info.subquery_source_tables,
                dialect,
                &mut sources,
            );
        } else {
            // ValuesOnly extracts value-producing refs (CASE conditions + results,
            // function args) but skips window PARTITION/ORDER BY.
            // Fallback to All only for truly empty (all-literal) expressions.
            let mut raw_refs = Vec::new();
            extract_columns_from_expr_values_only(expr, &mut raw_refs);
            if raw_refs.is_empty() {
                extract_all_columns_from_expr(expr, &mut raw_refs);
            }
            for raw in &raw_refs {
                // Use CTE-dict-aware resolver (much more powerful for CTE/subquery refs)
                let resolved = resolve_ref_with_cte_dict(
                    raw,
                    &info.alias_map,
                    flat_schema,
                    cte_dict,
                    dialect,
                    None,
                );
                sources.extend(resolved);
            }

            // Handle functions with wildcard args (e.g., OBJECT_CONSTRUCT(*)):
            // expand * to all columns from all FROM-clause tables in scope.
            if expr_contains_wildcard_func_arg(expr) {
                expand_wildcard_func_sources(&mut sources, info, flat_schema, cte_dict, dialect);
            }
        }

        if !sources.is_empty() {
            column_dict.entry(col_name).or_default().extend(sources);
        }
    }
}

/// Functions where `*` means "all columns" (data-consuming wildcards).
/// `COUNT(*)` just counts rows — it doesn't reference column data.
const WILDCARD_DATA_FUNCTIONS: &[&str] = &[
    "OBJECT_CONSTRUCT",
    "OBJECT_AGG",
    "TO_JSON",
    "TO_VARIANT",
    "HASH",
    "HASH_AGG",
];

/// Check if an expression tree contains a data-consuming function with wildcard.
fn expr_contains_wildcard_func_arg(expr: &Expr) -> bool {
    match expr {
        Expr::Function(f) => {
            let func_name = f
                .name
                .0
                .last()
                .map(|i| i.value.to_uppercase())
                .unwrap_or_default();
            if WILDCARD_DATA_FUNCTIONS.contains(&func_name.as_str())
                && let ast::FunctionArguments::List(arg_list) = &f.args
            {
                for arg in &arg_list.args {
                    if matches!(
                        arg,
                        ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Wildcard)
                    ) {
                        return true;
                    }
                }
            }
            // Recurse into nested function args
            if let ast::FunctionArguments::List(arg_list) = &f.args {
                for arg in &arg_list.args {
                    match arg {
                        ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e))
                        | ast::FunctionArg::Named {
                            arg: ast::FunctionArgExpr::Expr(e),
                            ..
                        } => {
                            if expr_contains_wildcard_func_arg(e) {
                                return true;
                            }
                        }
                        _ => {}
                    }
                }
            }
            false
        }
        Expr::Cast { expr: e, .. } | Expr::Nested(e) => expr_contains_wildcard_func_arg(e),
        _ => false,
    }
}

/// Expand wildcard function args to all columns from FROM-clause tables.
fn expand_wildcard_func_sources(
    sources: &mut BTreeSet<String>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    for actual_table in info.alias_map.values() {
        let actual = actual_table.clone();
        // Check if it's a CTE
        if let Some(cte_cols) = find_cte_in_dict(cte_dict, &actual) {
            for col_sources in cte_cols.values() {
                sources.extend(col_sources.iter().cloned());
            }
            continue;
        }
        // Check if it's a schema table
        if let Some(canonical) = find_schema_table(flat_schema, &actual)
            && let Some(cols) = flat_schema.get(canonical)
        {
            for col in cols {
                sources.insert(qualify_ref(canonical, col, dialect));
            }
        }
    }
}

/// Supplement existing column_dict entries with additional sources from the CTE dict.
///
/// After polyglot edge processing and star expansion, some output columns may
/// have partial source sets (e.g., only the leftmost UNION branch for SELECT *
/// over a UNION subquery). The CTE dict, built from our AST walking, has the
/// complete picture (all UNION branches, all CTE pass-throughs). This function
/// merges those additional sources into existing column_dict entries.
///
/// Expand inner tables (from subquery source tracking) into column_dict.
///
/// For each inner table, checks CTE dict first, then schema. Adds matching
/// column sources to existing column_dict entries.
fn expand_inner_tables_to_column_dict(
    inner_tables: &[String],
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    for inner_table in inner_tables {
        if let Some(cte_cols) = find_cte_in_dict(cte_dict, inner_table) {
            for (cte_col, cte_sources) in cte_cols {
                let col_key = normalize_column_name(cte_col, dialect);
                if let Some(existing) = column_dict.get_mut(&col_key) {
                    for src in cte_sources {
                        existing.insert(src.clone());
                    }
                }
            }
        } else if let Some(schema_name) = find_schema_table(flat_schema, inner_table)
            && let Some(cols) = flat_schema.get(schema_name)
        {
            for col in cols {
                let col_key = normalize_column_name(col, dialect);
                if let Some(existing) = column_dict.get_mut(&col_key) {
                    existing.insert(qualify_ref(schema_name, col, dialect));
                }
            }
        }
    }
}

/// Supplement column_dict with sources from non-leftmost UNION branches.
///
/// Polyglot-sql only generates edges from the leftmost UNION branch. This
/// function walks all branches of a set operation, resolves each SELECT item's
/// expression to find additional source refs, and merges them into column_dict.
fn supplement_union_branches(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    body: &SetExpr,
    outer_select_items: &[SelectItem],
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    // Collect all non-leftmost UNION branch bodies
    let mut branches = Vec::new();
    collect_union_branches(body, &mut branches, false);

    if branches.is_empty() {
        return;
    }

    for branch_body in branches {
        if let SetExpr::Select(sel) = branch_body {
            let mut sub = BTreeMap::new();
            let alias_map = build_alias_map(&sel.from, &mut sub);

            // Check if ALL branch items are wildcards (star-to-star UNION)
            let branch_all_stars = sel.projection.iter().all(|p| {
                matches!(
                    p,
                    SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(..)
                )
            });

            // For star-to-star UNIONs (e.g., SELECT * FROM t1 UNION ALL SELECT * FROM t2),
            // expand branch source tables into matching column_dict entries by column name.
            if branch_all_stars {
                for (alias_key, actual_table) in &alias_map {
                    if actual_table == "__subquery__" {
                        if let Some(inner_tables) = sub.get(alias_key) {
                            expand_inner_tables_to_column_dict(
                                inner_tables,
                                column_dict,
                                flat_schema,
                                cte_dict,
                                dialect,
                            );
                        }
                        continue;
                    }
                    if let Some(schema_name) = find_schema_table(flat_schema, actual_table) {
                        if let Some(cols) = flat_schema.get(schema_name) {
                            for col in cols {
                                let col_key = normalize_column_name(col, dialect);
                                if let Some(existing) = column_dict.get_mut(&col_key) {
                                    existing.insert(qualify_ref(schema_name, col, dialect));
                                }
                            }
                        }
                    } else if let Some(cte_cols) = find_cte_in_dict(cte_dict, actual_table) {
                        for (cte_col, cte_sources) in cte_cols {
                            let col_key = normalize_column_name(cte_col, dialect);
                            if let Some(existing) = column_dict.get_mut(&col_key) {
                                for src in cte_sources {
                                    existing.insert(src.clone());
                                }
                            }
                        }
                    }
                }
                continue;
            }

            // Named items: match branch SELECT items to outer output columns positionally.
            // Handle wildcards by expanding to column_dict keys from the first branch.
            for (i, item) in sel.projection.iter().enumerate() {
                // Handle wildcards: for each existing column_dict entry, look up the
                // matching column in the branch's source tables.
                if matches!(
                    item,
                    SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(..)
                ) {
                    // Get the table alias for qualified wildcards (e.g., t1.*)
                    let wildcard_table = match item {
                        SelectItem::QualifiedWildcard(name, _) => {
                            let parts: Vec<_> = name
                                .0
                                .iter()
                                .map(|p| normalize_column_name(&p.value, dialect))
                                .collect();
                            parts.last().cloned()
                        }
                        _ => None,
                    };

                    // Find the resolved table name for this wildcard
                    let resolved_table = wildcard_table.as_ref().and_then(|alias| {
                        alias_map
                            .get(&alias.to_uppercase())
                            .or_else(|| alias_map.get(alias))
                    });

                    if let Some(actual_table) = resolved_table {
                        // Schema lookup: add columns from this table for matching keys
                        if let Some(schema_name) = find_schema_table(flat_schema, actual_table) {
                            if let Some(cols) = flat_schema.get(schema_name) {
                                for col in cols {
                                    let col_key = normalize_column_name(col, dialect);
                                    if let Some(existing) = column_dict.get_mut(&col_key) {
                                        existing.insert(qualify_ref(schema_name, col, dialect));
                                    }
                                }
                            }
                        } else if let Some(cte_cols) = find_cte_in_dict(cte_dict, actual_table) {
                            for (cte_col, cte_sources) in cte_cols {
                                let col_key = normalize_column_name(cte_col, dialect);
                                if let Some(existing) = column_dict.get_mut(&col_key) {
                                    for src in cte_sources {
                                        existing.insert(src.clone());
                                    }
                                }
                            }
                        } else {
                            // No schema/CTE: mirror column_dict entries with this table
                            let table_norm = normalize_column_name(actual_table, dialect);
                            let entries: Vec<(String, String)> = column_dict
                                .keys()
                                .map(|k| (k.clone(), qualify_ref(&table_norm, k, dialect)))
                                .collect();
                            for (col_key, ref_str) in entries {
                                if let Some(existing) = column_dict.get_mut(&col_key) {
                                    existing.insert(ref_str);
                                }
                            }
                        }
                    } else {
                        // Unqualified wildcard: add all alias_map tables
                        for actual_table in alias_map.values() {
                            if actual_table == "__subquery__" {
                                continue;
                            }
                            if let Some(schema_name) = find_schema_table(flat_schema, actual_table)
                            {
                                if let Some(cols) = flat_schema.get(schema_name) {
                                    for col in cols {
                                        let col_key = normalize_column_name(col, dialect);
                                        if let Some(existing) = column_dict.get_mut(&col_key) {
                                            existing.insert(qualify_ref(schema_name, col, dialect));
                                        }
                                    }
                                }
                            } else if let Some(cte_cols) = find_cte_in_dict(cte_dict, actual_table)
                            {
                                for (cte_col, cte_sources) in cte_cols {
                                    let col_key = normalize_column_name(cte_col, dialect);
                                    if let Some(existing) = column_dict.get_mut(&col_key) {
                                        for src in cte_sources {
                                            existing.insert(src.clone());
                                        }
                                    }
                                }
                            }
                        }
                    }
                    continue;
                }

                let col_name = if i < outer_select_items.len() {
                    select_item_name(&outer_select_items[i], i, dialect)
                } else {
                    select_item_name(item, i, dialect)
                };
                let col_name = match col_name {
                    Some(n) if n != "*" && !n.ends_with(".*") => n,
                    _ => continue,
                };

                if !column_dict.contains_key(&col_name) {
                    continue;
                }

                let expr = match item {
                    SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => e,
                    _ => continue,
                };

                let mut raw_refs = Vec::new();
                extract_columns_from_expr(expr, &mut raw_refs);
                for raw in &raw_refs {
                    let resolved = resolve_ref_with_cte_dict(
                        raw,
                        &alias_map,
                        flat_schema,
                        cte_dict,
                        dialect,
                        None,
                    );
                    for r in &resolved {
                        column_dict
                            .entry(col_name.clone())
                            .or_default()
                            .insert(r.clone());
                    }
                }
            }
        }
    }
}

/// Collect non-leftmost bodies from a UNION set operation.
/// Only collects from UNION / UNION ALL — skips EXCEPT, INTERSECT, MINUS
/// (those filter rows from the right side, not merge data).
fn collect_union_branches<'a>(
    body: &'a SetExpr,
    branches: &mut Vec<&'a SetExpr>,
    skip_leftmost: bool,
) {
    match body {
        SetExpr::SetOperation {
            left, right, op, ..
        } => {
            if matches!(op, ast::SetOperator::Union) {
                collect_union_branches(left, branches, skip_leftmost);
                collect_union_branches(right, branches, true);
            }
            // For EXCEPT/INTERSECT/MINUS, don't collect right-side branches
        }
        SetExpr::Select(_) => {
            if skip_leftmost {
                branches.push(body);
            }
        }
        SetExpr::Query(q) => {
            collect_union_branches(&q.body, branches, skip_leftmost);
        }
        _ => {}
    }
}

/// Collect ALL branch bodies from a UNION set operation (including leftmost).
fn collect_all_union_branches<'a>(body: &'a SetExpr, branches: &mut Vec<&'a SetExpr>) {
    match body {
        SetExpr::SetOperation {
            left, right, op, ..
        } => {
            if matches!(op, ast::SetOperator::Union) {
                collect_all_union_branches(left, branches);
                collect_all_union_branches(right, branches);
            }
        }
        SetExpr::Select(_) => {
            branches.push(body);
        }
        SetExpr::Query(q) => {
            collect_all_union_branches(&q.body, branches);
        }
        _ => {}
    }
}

/// Supplement column_dict with sources from UNION branches inside FROM subqueries.
///
/// Polyglot-sql traces column lineage inconsistently through UNION branches:
/// bare column references (e.g., `member_id`) trace through all branches, but
/// aggregate-wrapped columns (e.g., `MAX(click_date) AS date`) often only trace
/// through one branch. This function finds all FROM-clause subqueries whose body
/// is a UNION, walks ALL branches, resolves column refs, and merges into column_dict.
fn supplement_from_subquery_unions(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    body: &SetExpr,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    if let SetExpr::Select(sel) = body {
        for twj in &sel.from {
            supplement_union_in_factor(&twj.relation, column_dict, flat_schema, cte_dict, dialect);
            for join in &twj.joins {
                supplement_union_in_factor(
                    &join.relation,
                    column_dict,
                    flat_schema,
                    cte_dict,
                    dialect,
                );
            }
        }
    }
}

/// Process a single FROM table factor, looking for UNION subqueries.
fn supplement_union_in_factor(
    factor: &ast::TableFactor,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    let subquery = match factor {
        ast::TableFactor::Derived { subquery, .. } => subquery,
        _ => return,
    };

    // Check if this subquery's body is a UNION (directly or through nested subqueries)
    let union_body = find_union_body(&subquery.body);
    let Some(union_body) = union_body else {
        // Not a UNION — but recurse into the subquery's FROM to find nested ones
        supplement_from_subquery_unions(
            column_dict,
            &subquery.body,
            flat_schema,
            cte_dict,
            dialect,
        );
        return;
    };

    // Collect ALL branches
    let mut branches = Vec::new();
    collect_all_union_branches(union_body, &mut branches);

    // Determine canonical column names from the first branch.
    // In SQL UNIONs, column names come from the FIRST branch — subsequent
    // branches contribute data by POSITION, not by name.
    let first_branch_names: Vec<String> = branches
        .first()
        .map(|b| {
            collect_branch_named_items(b, flat_schema, cte_dict, dialect, 0)
                .into_iter()
                .map(|(name, _)| name)
                .collect()
        })
        .unwrap_or_default();

    for branch_body in &branches {
        let items = collect_branch_named_items(branch_body, flat_schema, cte_dict, dialect, 0);
        for (pos, (branch_name, sources)) in items.iter().enumerate() {
            // When a UNION branch has a different alias than the first branch
            // at the same position, use the first branch's name (SQL standard:
            // UNION column names come from the first branch). This prevents
            // accidental name collisions where a branch alias matches an outer
            // column name it shouldn't map to.
            let col_name = match first_branch_names.get(pos) {
                Some(first_name) if first_name != branch_name => first_name,
                _ => branch_name,
            };
            if let Some(existing) = column_dict.get_mut(col_name) {
                for s in sources {
                    existing.insert(s.clone());
                }
            }
        }

        // Recurse into this branch's FROM to find nested UNIONs
        if let SetExpr::Select(sel) = branch_body {
            for twj in &sel.from {
                supplement_union_in_factor(
                    &twj.relation,
                    column_dict,
                    flat_schema,
                    cte_dict,
                    dialect,
                );
                for join in &twj.joins {
                    supplement_union_in_factor(
                        &join.relation,
                        column_dict,
                        flat_schema,
                        cte_dict,
                        dialect,
                    );
                }
            }
        }
    }
}

/// Recursively collect named items from a SELECT body, walking through
/// `SELECT *` wrappers to find actual expressions with column references.
///
/// For `SELECT * FROM (SELECT a, MAX(b) AS max_b FROM t) sub`, this walks
/// through the `*` into the inner SELECT and returns `[(a, {t.a}), (max_b, {t.b})]`.
fn collect_branch_named_items(
    body: &SetExpr,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<(String, BTreeSet<String>)> {
    if depth > 10 {
        return vec![];
    }

    let sel = match body {
        SetExpr::Select(sel) => sel,
        SetExpr::Query(q) => {
            return collect_branch_named_items(&q.body, flat_schema, cte_dict, dialect, depth);
        }
        _ => return vec![],
    };

    let mut sub = BTreeMap::new();
    let alias_map = build_alias_map(&sel.from, &mut sub);
    let mut results = Vec::new();

    for (i, item) in sel.projection.iter().enumerate() {
        match item {
            SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(..) => {
                // Star item: recurse into FROM subqueries to find actual items
                for twj in &sel.from {
                    if let ast::TableFactor::Derived { subquery, .. } = &twj.relation {
                        results.extend(collect_branch_named_items(
                            &subquery.body,
                            flat_schema,
                            cte_dict,
                            dialect,
                            depth + 1,
                        ));
                    }
                }
                // Also add passthrough columns from base tables in scope
                for (alias_key, actual_table) in &alias_map {
                    if actual_table == "__subquery__" {
                        if let Some(inner_tables) = sub.get(alias_key) {
                            for inner_table in inner_tables {
                                add_table_columns_to_results(
                                    inner_table,
                                    flat_schema,
                                    cte_dict,
                                    dialect,
                                    &mut results,
                                );
                            }
                        }
                    } else {
                        add_table_columns_to_results(
                            actual_table,
                            flat_schema,
                            cte_dict,
                            dialect,
                            &mut results,
                        );
                    }
                }
            }
            _ => {
                // Named item: resolve its expression
                let col_name = match select_item_name(item, i, dialect) {
                    Some(n) if n != "*" && !n.ends_with(".*") => n,
                    _ => continue,
                };

                let expr = match item {
                    SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => e,
                    _ => continue,
                };

                let mut sources = BTreeSet::new();
                let mut raw_refs = Vec::new();
                extract_columns_from_expr(expr, &mut raw_refs);
                for raw in &raw_refs {
                    let resolved = resolve_ref_with_cte_dict(
                        raw,
                        &alias_map,
                        flat_schema,
                        cte_dict,
                        dialect,
                        None,
                    );
                    for r in &resolved {
                        sources.insert(r.clone());
                    }
                }
                if !sources.is_empty() {
                    results.push((col_name, sources));
                }
            }
        }
    }

    results
}

/// Add columns from a table (by name) to results as passthrough entries.
fn add_table_columns_to_results(
    table_name: &str,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    results: &mut Vec<(String, BTreeSet<String>)>,
) {
    if let Some(cte_cols) = find_cte_in_dict(cte_dict, table_name) {
        for (cte_col, cte_sources) in cte_cols {
            let col_key = normalize_column_name(cte_col, dialect);
            results.push((col_key, cte_sources.clone()));
        }
    } else if let Some(schema_name) = find_schema_table(flat_schema, table_name)
        && let Some(cols) = flat_schema.get(schema_name)
    {
        for col in cols {
            let col_key = normalize_column_name(col, dialect);
            let source = qualify_ref(schema_name, col, dialect);
            results.push((col_key, BTreeSet::from([source])));
        }
    }
}

/// Find the innermost UNION body in a SetExpr, walking through wrapping Selects.
/// Returns the SetOperation body if found, None otherwise.
fn find_union_body(body: &SetExpr) -> Option<&SetExpr> {
    match body {
        SetExpr::SetOperation { op, .. } => {
            if matches!(op, ast::SetOperator::Union) {
                Some(body)
            } else {
                None
            }
        }
        SetExpr::Select(sel) => {
            // SELECT * FROM (... UNION ALL ...) — look through single FROM subqueries
            if sel.from.len() == 1
                && sel.from[0].joins.is_empty()
                && let ast::TableFactor::Derived { subquery, .. } = &sel.from[0].relation
            {
                return find_union_body(&subquery.body);
            }
            None
        }
        SetExpr::Query(q) => find_union_body(&q.body),
        _ => None,
    }
}

/// Remap SELECT output columns to INSERT INTO target columns positionally.
///
/// For `INSERT INTO t (col_a, col_b) SELECT x, y, z FROM ...`, the SELECT
/// produces columns x, y, z but the INSERT target only has col_a and col_b.
/// This function:
/// 1. Builds a positional mapping: SELECT item 0 → INSERT col 0, etc.
/// 2. Renames column_dict keys from SELECT names to INSERT target names
/// 3. Discards extra SELECT columns beyond the INSERT column list
fn remap_insert_columns(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    _flat_schema: &FlatSchema,
    dialect: SqlDialect,
) {
    let target_cols = &info.insert_target_columns;
    if target_cols.is_empty() {
        return;
    }

    // Build ordered SELECT column names from the select_items
    let select_names: Vec<String> = info
        .select_items
        .iter()
        .enumerate()
        .filter_map(|(i, item)| select_item_name(item, i, dialect))
        .filter(|name| name != "*" && !name.ends_with(".*"))
        .collect();

    // If SELECT has fewer named columns than target, we can't remap reliably
    if select_names.is_empty() {
        return;
    }

    // Build the old→new name mapping (positional)
    let mut remap: HashMap<String, String> = HashMap::new();
    let mut keep_set: BTreeSet<String> = BTreeSet::new();
    for (i, target_col) in target_cols.iter().enumerate() {
        let target_normalized = normalize_column_name(target_col, dialect);
        keep_set.insert(target_normalized.clone());

        if i < select_names.len() {
            let select_name = &select_names[i];
            if *select_name != target_normalized {
                remap.insert(select_name.clone(), target_normalized);
            }
        }
    }

    // If the target columns match the current keys with no remapping needed
    // and there are no extras, nothing to do
    if remap.is_empty() {
        // Just remove extra keys not in the target list
        let extras: Vec<String> = column_dict
            .keys()
            .filter(|k| !keep_set.contains(k.as_str()))
            .cloned()
            .collect();
        for k in extras {
            column_dict.remove(&k);
        }
        return;
    }

    // Apply remapping: rename keys and remove extras
    let mut new_dict: HashMap<String, BTreeSet<String>> = HashMap::new();

    // First, handle columns that need remapping (by position)
    for (i, target_col) in target_cols.iter().enumerate() {
        let target_normalized = normalize_column_name(target_col, dialect);
        if i < select_names.len() {
            let select_name = &select_names[i];
            if let Some(sources) = column_dict.get(select_name) {
                new_dict
                    .entry(target_normalized)
                    .or_default()
                    .extend(sources.iter().cloned());
            }
        }
    }

    // For columns already in target list that weren't remapped (e.g., from star expansion
    // that already produced the right name), keep them
    for target_col in target_cols {
        let target_normalized = normalize_column_name(target_col, dialect);
        if !new_dict.contains_key(&target_normalized)
            && let Some(sources) = column_dict.get(&target_normalized)
        {
            new_dict.insert(target_normalized, sources.clone());
        }
    }

    // Also check: if star expansion produced columns that match target names
    // but aren't in select_names, include them
    for target_col in target_cols {
        let target_normalized = normalize_column_name(target_col, dialect);
        if !new_dict.contains_key(&target_normalized) {
            // Try case-insensitive lookup in column_dict
            let lower = target_normalized.to_lowercase();
            for (k, v) in column_dict.iter() {
                if k.to_lowercase() == lower {
                    new_dict.insert(target_normalized.clone(), v.clone());
                    break;
                }
            }
        }
    }

    *column_dict = new_dict;
}

/// When INSERT INTO has no explicit column list, try to infer target columns
/// from the schema and use those to filter extra output columns.
fn infer_insert_target_columns(
    stmts: &[Statement],
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) -> Vec<String> {
    for stmt in stmts {
        if let Statement::Insert(ins) = stmt {
            if !ins.columns.is_empty() {
                // Explicit column list already handled in analyze_query
                return vec![];
            }
            // Extract target table name
            let target_table = match &ins.table {
                ast::TableObject::TableName(name) => name.to_string(),
                _ => continue,
            };
            let target_table_clean = target_table
                .split('.')
                .next_back()
                .unwrap_or(&target_table)
                .trim_matches('"')
                .to_string();

            // Look up in schema
            if let Some(canonical) = find_schema_table(flat_schema, &target_table_clean)
                && let Some(cols) = flat_schema.get(canonical)
            {
                return cols
                    .iter()
                    .map(|c| normalize_column_name(c, dialect))
                    .collect();
            }
        }
    }
    vec![]
}

/// Resolve columns from a scalar subquery using its own alias map.
fn resolve_scalar_subquery_columns(
    subquery: &ast::Query,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
    sources: &mut BTreeSet<String>,
) {
    let mut sub = BTreeMap::new();
    let subq_alias = build_alias_map_from_body(&subquery.body, &mut sub);

    // Extract ALL column refs from SELECT + WHERE + HAVING + GROUP BY
    let mut raw_refs = Vec::new();
    extract_all_columns_from_body(&subquery.body, &mut raw_refs);

    for raw in &raw_refs {
        let resolved = resolve_column_ref(
            raw,
            &subq_alias,
            flat_schema,
            cte_names,
            cte_source_tables,
            subquery_sources,
            dialect,
        );
        sources.extend(resolved);
    }
}

/// Extract ALL column refs from an expression, INCLUDING from subquery SELECT lists.
/// Used for expression column resolution (scalar subqueries, etc).
fn extract_all_columns_from_expr(expr: &Expr, refs: &mut Vec<RawColumnRef>) {
    match expr {
        Expr::Subquery(subquery) => {
            // For scalar subqueries: extract from both SELECT and WHERE
            extract_all_columns_from_body(&subquery.body, refs);
        }
        Expr::InSubquery {
            expr: e, subquery, ..
        } => {
            extract_columns_from_expr(e, refs);
            extract_all_columns_from_body(&subquery.body, refs);
        }
        Expr::Exists { subquery, .. } => {
            extract_all_columns_from_body(&subquery.body, refs);
        }
        // For non-subquery expressions, use the regular extractor
        other => {
            extract_columns_from_expr(other, refs);
        }
    }
}

fn extract_all_columns_from_body(body: &SetExpr, refs: &mut Vec<RawColumnRef>) {
    match body {
        SetExpr::Select(sel) => {
            // SELECT list
            for item in &sel.projection {
                match item {
                    SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
                        extract_columns_from_expr(e, refs);
                    }
                    _ => {}
                }
            }
            // WHERE
            if let Some(w) = &sel.selection {
                extract_columns_from_expr(w, refs);
            }
            // HAVING
            if let Some(h) = &sel.having {
                extract_columns_from_expr(h, refs);
            }
            // GROUP BY
            if let GroupByExpr::Expressions(exprs, _) = &sel.group_by {
                for expr in exprs {
                    if !is_numeric_literal(expr) {
                        extract_columns_from_expr(expr, refs);
                    }
                }
            }
            // JOIN ON
            for twj in &sel.from {
                for join in &twj.joins {
                    match &join.join_operator {
                        JoinOperator::Inner(c)
                        | JoinOperator::LeftOuter(c)
                        | JoinOperator::RightOuter(c)
                        | JoinOperator::FullOuter(c) => {
                            if let JoinConstraint::On(expr) = c {
                                extract_columns_from_expr(expr, refs);
                            }
                        }
                        _ => {}
                    }
                }
            }
        }
        SetExpr::SetOperation { left, right, .. } => {
            extract_all_columns_from_body(left, refs);
            extract_all_columns_from_body(right, refs);
        }
        SetExpr::Query(q) => extract_all_columns_from_body(&q.body, refs),
        _ => {}
    }
}

// ===========================================================================
// Star expansion
// ===========================================================================

fn expand_stars(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    let star_key = normalize_column_name("*", dialect);
    // Only expand unqualified * if the outer SELECT actually has SELECT *
    let has_unqualified_star = info
        .select_items
        .iter()
        .any(|i| matches!(i, SelectItem::Wildcard(_)));

    // Collect qualified wildcard prefixes from SELECT items (e.g., "t1" from "t1.*")
    let qualified_wildcard_prefixes: Vec<String> = info
        .select_items
        .iter()
        .filter_map(|item| {
            if let SelectItem::QualifiedWildcard(name, _) = item {
                Some(name.to_string())
            } else {
                None
            }
        })
        .collect();

    if let Some(star_sources) = column_dict.remove(&star_key) {
        if has_unqualified_star {
            let mut expanded_any = false;
            for source_ref in &star_sources {
                let parts: Vec<&str> = source_ref.split('.').collect();
                if parts.len() < 2 {
                    continue;
                }
                let table_parts: Vec<String> = parts[..parts.len() - 1]
                    .iter()
                    .map(|p| p.trim_matches('"').to_string())
                    .collect();
                let table_name = table_parts.join(".");

                // Skip empty table names (polyglot returns "" for subquery sources)
                if table_name.is_empty() {
                    continue;
                }

                // Check if this is a CTE (case-insensitive)
                let table_lower = table_name.to_lowercase();
                let is_cte = info
                    .cte_names
                    .iter()
                    .any(|c| c.to_lowercase() == table_lower);

                if is_cte {
                    expand_cte_star(
                        &table_name,
                        column_dict,
                        info,
                        flat_schema,
                        cte_dict,
                        dialect,
                    );
                    expanded_any = true;
                } else {
                    expand_table_star(&table_name, column_dict, info, flat_schema, dialect);
                    expanded_any = true;
                }
            }

            // If we couldn't expand any star sources, OR expansion produced nothing,
            // expand from the FROM clause — find all tables/subqueries in scope
            if !expanded_any || column_dict.is_empty() {
                expand_star_from_scope(column_dict, info, flat_schema, cte_dict, dialect);
            }
        } else if !qualified_wildcard_prefixes.is_empty() {
            // Outer SELECT has qualified wildcards (e.g., t1.*) but not unqualified *.
            // Polyglot returns target_column="*" for both cases, so we need to expand
            // only the tables matching the qualified wildcard prefixes.
            for prefix in &qualified_wildcard_prefixes {
                let actual_table = info
                    .alias_map
                    .iter()
                    .find(|(k, _)| k.to_lowercase() == prefix.to_lowercase())
                    .map(|(_, v)| v.clone())
                    .unwrap_or_else(|| prefix.clone());

                let is_cte = info
                    .cte_names
                    .iter()
                    .any(|c| c.to_lowercase() == actual_table.to_lowercase());

                if is_cte {
                    expand_cte_star(
                        &actual_table,
                        column_dict,
                        info,
                        flat_schema,
                        cte_dict,
                        dialect,
                    );
                } else if actual_table == "__subquery__" {
                    expand_subquery_star(prefix, column_dict, info, flat_schema, cte_dict, dialect);
                } else {
                    expand_table_star(&actual_table, column_dict, info, flat_schema, dialect);
                }
            }
        }
        // If neither unqualified nor qualified wildcard, drop the * entry — it's from internal CTEs
    } else if has_unqualified_star && column_dict.is_empty() {
        // Polyglot returned 0 edges (no * key in column_dict).
        // Proactively expand SELECT * from the AST scope using schema.
        expand_star_from_scope(column_dict, info, flat_schema, cte_dict, dialect);
    }

    // Also handle qualified wildcards like t1.* that polyglot might return
    let qualified_star_keys: Vec<String> = column_dict
        .keys()
        .filter(|k| k.ends_with(".*"))
        .cloned()
        .collect();
    for key in qualified_star_keys {
        if let Some(_sources) = column_dict.remove(&key) {
            let table_prefix = key.trim_end_matches(".*");
            // Case-insensitive alias lookup
            let actual_table = info
                .alias_map
                .iter()
                .find(|(k, _)| k.to_lowercase() == table_prefix.to_lowercase())
                .map(|(_, v)| v.clone())
                .unwrap_or_else(|| table_prefix.to_string());

            let is_cte = info
                .cte_names
                .iter()
                .any(|c| c.to_lowercase() == actual_table.to_lowercase());
            if is_cte {
                expand_cte_star(
                    &actual_table,
                    column_dict,
                    info,
                    flat_schema,
                    cte_dict,
                    dialect,
                );
            } else if actual_table == "__subquery__" {
                // Qualified wildcard on subquery: expand through subquery source tables
                if let Some(inner_tables) = info.subquery_source_tables.get(table_prefix) {
                    for inner_table in inner_tables {
                        let inner_lower = inner_table.to_lowercase();
                        let is_inner_cte = info
                            .cte_names
                            .iter()
                            .any(|c| c.to_lowercase() == inner_lower);
                        if is_inner_cte {
                            expand_cte_star(
                                inner_table,
                                column_dict,
                                info,
                                flat_schema,
                                cte_dict,
                                dialect,
                            );
                        } else {
                            expand_table_star(inner_table, column_dict, info, flat_schema, dialect);
                        }
                    }
                }
            } else {
                expand_table_star(&actual_table, column_dict, info, flat_schema, dialect);
            }
        }
    }
}

/// Expand `SELECT *` by looking at all tables in the FROM clause scope.
/// Used when polyglot returns empty source_table for star edges.
fn expand_star_from_scope(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    for item in &info.select_items {
        if matches!(item, SelectItem::Wildcard(_)) {
            for (alias, actual) in &info.alias_map {
                if actual == "__subquery__" {
                    // Subquery: check if it contains a UNION and use last branch only
                    expand_subquery_star(alias, column_dict, info, flat_schema, cte_dict, dialect);
                } else {
                    let is_cte = info
                        .cte_names
                        .iter()
                        .any(|c| c.to_lowercase() == actual.to_lowercase());
                    if is_cte {
                        expand_cte_star(actual, column_dict, info, flat_schema, cte_dict, dialect);
                    } else {
                        expand_table_star(actual, column_dict, info, flat_schema, dialect);
                    }
                }
            }
            break;
        }
    }
}

/// Expand star through a subquery alias, handling UNIONs correctly.
/// For UNION subqueries, trace lineage to ALL branches — each output column
/// depends on the corresponding column from every branch.
fn expand_subquery_star(
    alias: &str,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    // Find the subquery body in the parsed query to check for UNION
    let subquery_body = find_subquery_body(&info.query.body, alias);

    if let Some(body) = subquery_body {
        // Get tables from ALL branches of any set operation
        let tables = extract_all_union_tables(body);
        for table in &tables {
            let table_lower = table.to_lowercase();
            let is_cte = info
                .cte_names
                .iter()
                .any(|c| c.to_lowercase() == table_lower);
            if is_cte {
                expand_cte_star(table, column_dict, info, flat_schema, cte_dict, dialect);
            } else {
                expand_table_star(table, column_dict, info, flat_schema, dialect);
            }
        }
    } else if let Some(inner_tables) = info.subquery_source_tables.get(alias) {
        // Fallback: expand ALL tables from subquery_sources
        for table in inner_tables {
            let is_cte = info
                .cte_names
                .iter()
                .any(|c| c.to_lowercase() == table.to_lowercase());
            if is_cte {
                expand_cte_star(table, column_dict, info, flat_schema, cte_dict, dialect);
            } else {
                expand_table_star(table, column_dict, info, flat_schema, dialect);
            }
        }
    }
}

/// Recursively expand SELECT * through nested subqueries.
/// Handles patterns like `SELECT * FROM (SELECT * FROM (SELECT * FROM base_table))`.
fn expand_star_through_nested_subqueries(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    body: &SetExpr,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    depth: usize,
) {
    if depth > 10 {
        return;
    }
    match body {
        SetExpr::Select(sel) => {
            // Check if this SELECT has * and a FROM clause with subqueries/tables
            let has_star = sel
                .projection
                .iter()
                .any(|p| matches!(p, SelectItem::Wildcard(_)));
            if !has_star && depth > 0 {
                // This inner SELECT doesn't have *, so process its projection items
                for item in &sel.projection {
                    match item {
                        SelectItem::QualifiedWildcard(name, _) => {
                            let prefix = name.to_string();
                            let table_name = info
                                .alias_map
                                .iter()
                                .find(|(k, _)| k.eq_ignore_ascii_case(&prefix))
                                .map(|(_, v)| v.as_str())
                                .unwrap_or(&prefix);
                            let is_cte = info
                                .cte_names
                                .iter()
                                .any(|c| c.eq_ignore_ascii_case(table_name));
                            if is_cte {
                                expand_cte_star(
                                    table_name,
                                    column_dict,
                                    info,
                                    flat_schema,
                                    cte_dict,
                                    dialect,
                                );
                            } else {
                                expand_table_star(
                                    table_name,
                                    column_dict,
                                    info,
                                    flat_schema,
                                    dialect,
                                );
                            }
                        }
                        SelectItem::ExprWithAlias {
                            expr,
                            alias: col_alias,
                        } => {
                            let col_name = normalize_column_name(&col_alias.value, dialect);
                            let mut raw_refs = Vec::new();
                            extract_columns_from_expr(expr, &mut raw_refs);
                            for raw in &raw_refs {
                                let resolved = resolve_ref_with_cte_dict(
                                    raw,
                                    &info.alias_map,
                                    flat_schema,
                                    cte_dict,
                                    dialect,
                                    None,
                                );
                                for r in &resolved {
                                    column_dict
                                        .entry(col_name.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                            }
                        }
                        SelectItem::UnnamedExpr(expr) => {
                            let col_name = match expr {
                                Expr::Identifier(id) => normalize_column_name(&id.value, dialect),
                                Expr::CompoundIdentifier(parts) => normalize_column_name(
                                    &parts.last().map(|p| p.value.clone()).unwrap_or_default(),
                                    dialect,
                                ),
                                _ => continue,
                            };
                            let mut raw_refs = Vec::new();
                            extract_columns_from_expr(expr, &mut raw_refs);
                            for raw in &raw_refs {
                                let resolved = resolve_ref_with_cte_dict(
                                    raw,
                                    &info.alias_map,
                                    flat_schema,
                                    cte_dict,
                                    dialect,
                                    None,
                                );
                                for r in &resolved {
                                    column_dict
                                        .entry(col_name.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                            }
                        }
                        _ => {}
                    }
                }
                return;
            }

            // Has * — recurse into FROM clause
            for twj in &sel.from {
                expand_star_from_table_factor(
                    &twj.relation,
                    column_dict,
                    info,
                    flat_schema,
                    cte_dict,
                    dialect,
                    depth,
                );
                for join in &twj.joins {
                    expand_star_from_table_factor(
                        &join.relation,
                        column_dict,
                        info,
                        flat_schema,
                        cte_dict,
                        dialect,
                        depth,
                    );
                }
            }
        }
        SetExpr::SetOperation { left, .. } => {
            // For UNION/INTERSECT/EXCEPT, expand from the first branch
            expand_star_through_nested_subqueries(
                column_dict,
                left,
                info,
                flat_schema,
                cte_dict,
                dialect,
                depth + 1,
            );
        }
        SetExpr::Query(q) => {
            expand_star_through_nested_subqueries(
                column_dict,
                &q.body,
                info,
                flat_schema,
                cte_dict,
                dialect,
                depth + 1,
            );
        }
        _ => {}
    }
}

/// Helper for expand_star_through_nested_subqueries: process a table factor.
fn expand_star_from_table_factor(
    factor: &ast::TableFactor,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    depth: usize,
) {
    match factor {
        ast::TableFactor::Derived { subquery, .. } => {
            // Recurse into subquery
            expand_star_through_nested_subqueries(
                column_dict,
                &subquery.body,
                info,
                flat_schema,
                cte_dict,
                dialect,
                depth + 1,
            );
        }
        ast::TableFactor::Table { name, .. } => {
            let table_name = name
                .0
                .iter()
                .map(|p| p.value.clone())
                .collect::<Vec<_>>()
                .join(".");
            let is_cte = info
                .cte_names
                .iter()
                .any(|c| c.eq_ignore_ascii_case(&table_name));
            if is_cte {
                expand_cte_star(
                    &table_name,
                    column_dict,
                    info,
                    flat_schema,
                    cte_dict,
                    dialect,
                );
            } else {
                expand_table_star(&table_name, column_dict, info, flat_schema, dialect);
            }
        }
        ast::TableFactor::NestedJoin {
            table_with_joins, ..
        } => {
            expand_star_from_table_factor(
                &table_with_joins.relation,
                column_dict,
                info,
                flat_schema,
                cte_dict,
                dialect,
                depth,
            );
            for join in &table_with_joins.joins {
                expand_star_from_table_factor(
                    &join.relation,
                    column_dict,
                    info,
                    flat_schema,
                    cte_dict,
                    dialect,
                    depth,
                );
            }
        }
        _ => {}
    }
}

/// Find the subquery body for a given alias in the FROM clause.
fn find_subquery_body<'a>(body: &'a SetExpr, alias: &str) -> Option<&'a SetExpr> {
    match body {
        SetExpr::Select(sel) => {
            for twj in &sel.from {
                if let Some(b) = find_subquery_body_in_factor(&twj.relation, alias) {
                    return Some(b);
                }
                for join in &twj.joins {
                    if let Some(b) = find_subquery_body_in_factor(&join.relation, alias) {
                        return Some(b);
                    }
                }
            }
            None
        }
        _ => None,
    }
}

fn find_subquery_body_in_factor<'a>(
    factor: &'a ast::TableFactor,
    alias: &str,
) -> Option<&'a SetExpr> {
    if let ast::TableFactor::Derived {
        alias: Some(a),
        subquery,
        ..
    } = factor
        && a.name.value.to_lowercase() == alias.to_lowercase()
    {
        return Some(&subquery.body);
    }
    None
}

/// Extract table names from the rightmost branch of a set operation.
/// For non-set-operations, returns all tables from the SELECT.
/// Extract tables from ALL branches of a UNION set operation.
fn extract_all_union_tables(body: &SetExpr) -> Vec<String> {
    match body {
        SetExpr::SetOperation { left, right, .. } => {
            let mut tables = extract_all_union_tables(left);
            tables.extend(extract_all_union_tables(right));
            tables
        }
        SetExpr::Select(sel) => {
            let mut tables = Vec::new();
            for twj in &sel.from {
                extract_table_name_from_factor(&twj.relation, &mut tables);
                for join in &twj.joins {
                    extract_table_name_from_factor(&join.relation, &mut tables);
                }
            }
            tables
        }
        SetExpr::Query(q) => extract_all_union_tables(&q.body),
        _ => vec![],
    }
}

fn expand_table_star(
    table_name: &str,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) {
    expand_table_star_inner(table_name, column_dict, info, flat_schema, dialect, true);
}

/// Inner implementation with `apply_set_op_filter` flag. When called from
/// `expand_cte_star`, set-operation filtering already happened at the CTE edge
/// level, so we skip it here to avoid comparing base table names against CTE aliases.
fn expand_table_star_inner(
    table_name: &str,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    apply_set_op_filter: bool,
) {
    // Use find_schema_table for multi-part name resolution
    let matched_table = find_schema_table(flat_schema, table_name);
    let table_lower = table_name.to_lowercase();
    for (schema_table, columns) in flat_schema.iter() {
        let is_match = match &matched_table {
            Some(mt) => schema_table.as_str() == *mt,
            None => schema_table.to_lowercase() == table_lower,
        };
        if is_match {
            if apply_set_op_filter && info.is_set_operation && !info.left_side_tables.is_empty() {
                let is_left = info
                    .left_side_tables
                    .iter()
                    .any(|t| t.to_lowercase() == table_lower);
                if !is_left {
                    continue;
                }
            }
            for col in columns {
                let col_name = normalize_column_name(col, dialect);
                let source = qualify_ref(schema_table, col, dialect);
                column_dict.entry(col_name).or_default().insert(source);
            }
            break;
        }
    }
}

fn expand_cte_star(
    cte_name: &str,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    expand_cte_star_inner(
        cte_name,
        column_dict,
        info,
        flat_schema,
        cte_dict,
        dialect,
        0,
    );
}

fn expand_cte_star_inner(
    cte_name: &str,
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
    depth: usize,
) {
    if depth > 10 {
        return;
    }

    // Try CTE dict first — if we have a pre-computed column map, use it directly
    if let Some(cte_cols) = find_cte_in_dict(cte_dict, cte_name) {
        for (col_name, sources) in cte_cols {
            for r in sources {
                column_dict
                    .entry(col_name.clone())
                    .or_default()
                    .insert(r.clone());
            }
        }
        return;
    }

    // Fallback to existing approach when CTE dict doesn't have this CTE
    let cte_lower = cte_name.to_lowercase();
    let output_cols = info
        .cte_output_columns
        .iter()
        .find(|(k, _)| k.to_lowercase() == cte_lower)
        .map(|(_, v)| v.clone())
        .unwrap_or_default();

    for col in &output_cols {
        if col == "*" {
            // CTE itself has *, expand through its source tables.
            // If source is another CTE, recurse; otherwise expand from schema.
            let mut handled = false;
            if let Some(source_tables) = info.cte_source_tables.get(cte_name) {
                let tables_clone: Vec<String> = source_tables.clone();
                for src_table in &tables_clone {
                    let is_inner_cte = info
                        .cte_names
                        .iter()
                        .any(|c| c.eq_ignore_ascii_case(src_table));
                    if is_inner_cte {
                        expand_cte_star_inner(
                            src_table,
                            column_dict,
                            info,
                            flat_schema,
                            cte_dict,
                            dialect,
                            depth + 1,
                        );
                    } else {
                        expand_table_star_inner(
                            src_table,
                            column_dict,
                            info,
                            flat_schema,
                            dialect,
                            false,
                        );
                    }
                }
                handled = true;
            }
            // Also check case-insensitive CTE name lookup
            if !handled {
                let cte_lower_key = cte_name.to_lowercase();
                for (k, source_tables) in &info.cte_source_tables {
                    if k.to_lowercase() == cte_lower_key {
                        let tables_clone: Vec<String> = source_tables.clone();
                        for src_table in &tables_clone {
                            let is_inner_cte = info
                                .cte_names
                                .iter()
                                .any(|c| c.eq_ignore_ascii_case(src_table));
                            if is_inner_cte {
                                expand_cte_star_inner(
                                    src_table,
                                    column_dict,
                                    info,
                                    flat_schema,
                                    cte_dict,
                                    dialect,
                                    depth + 1,
                                );
                            } else {
                                expand_table_star_inner(
                                    src_table,
                                    column_dict,
                                    info,
                                    flat_schema,
                                    dialect,
                                    false,
                                );
                            }
                        }
                    }
                }
            }
        } else {
            let col_name = normalize_column_name(col, dialect);
            let resolved =
                resolve_cte_column(cte_name, col, &info.cte_source_tables, flat_schema, dialect);
            if !resolved.is_empty() {
                for r in &resolved {
                    column_dict
                        .entry(col_name.clone())
                        .or_default()
                        .insert(r.clone());
                }
            } else {
                // Column is an expression alias in the CTE (e.g., TO_TIMESTAMP(col) AS event_ts).
                // Trace through the CTE's SELECT expression to find base table dependencies.
                let expr_deps = resolve_cte_expression_column(
                    cte_name,
                    col,
                    &info.query,
                    flat_schema,
                    &info.cte_names,
                    &info.cte_source_tables,
                    &info.subquery_source_tables,
                    dialect,
                );
                for dep in &expr_deps {
                    column_dict
                        .entry(col_name.clone())
                        .or_default()
                        .insert(dep.clone());
                }
            }
        }
    }
}

// ===========================================================================
// CTE expression column resolution
//
// When a CTE output column doesn't map to a base table column (e.g.,
// RANK() OVER ...), we need to look at the CTE's SELECT list to find
// the defining expression and extract its column dependencies.
// ===========================================================================

#[allow(clippy::too_many_arguments)]
fn resolve_cte_expression_column(
    cte_name: &str,
    column: &str,
    query: &ast::Query,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> Vec<String> {
    let cte_lower = cte_name.to_lowercase();
    let col_lower = column.to_lowercase();

    if let Some(with) = &query.with {
        for cte in &with.cte_tables {
            if cte.alias.name.value.to_lowercase() != cte_lower {
                continue;
            }

            // Find the SELECT item that defines this column
            let items = extract_leftmost_select_items(&cte.query.body);
            let mut sub = BTreeMap::new();
            let cte_alias_map = build_alias_map_from_body(&cte.query.body, &mut sub);

            for item in &items {
                let item_name = match item {
                    SelectItem::ExprWithAlias { alias, .. } => alias.value.clone(),
                    SelectItem::UnnamedExpr(Expr::Identifier(ident)) => ident.value.clone(),
                    SelectItem::UnnamedExpr(Expr::CompoundIdentifier(parts)) => {
                        parts.last().map(|p| p.value.clone()).unwrap_or_default()
                    }
                    _ => continue,
                };

                if item_name.to_lowercase() != col_lower {
                    continue;
                }

                // Extract ALL column refs from this expression
                let mut raw_refs = Vec::new();
                match item {
                    SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
                        extract_columns_from_expr(e, &mut raw_refs);
                    }
                    _ => {}
                }

                let mut resolved = Vec::new();
                for raw in &raw_refs {
                    let r = resolve_column_ref(
                        raw,
                        &cte_alias_map,
                        flat_schema,
                        cte_names,
                        cte_source_tables,
                        subquery_sources,
                        dialect,
                    );
                    resolved.extend(r);
                }

                return resolved;
            }
        }
    }

    vec![]
}

/// Resolve an expression alias through a subquery body.
///
/// Check if ALL qualified refs in the list point to tables in the schema.
fn all_refs_in_schema(refs: &[String], flat_schema: &FlatSchema) -> bool {
    if refs.is_empty() {
        return false;
    }
    refs.iter().all(|r| {
        // Parse "TABLE"."COL" — the table is the first segment
        let stripped = r.trim_matches('"');
        let dot = stripped.find('"');
        let table_part = if let Some(pos) = dot {
            &stripped[..pos]
        } else {
            let parts: Vec<&str> = stripped.split('.').collect();
            if parts.len() >= 2 {
                parts[0]
            } else {
                return false;
            }
        };
        let table = table_part.trim_matches('"');
        find_schema_table(flat_schema, table).is_some()
    })
}

/// When `resolve_subquery_column` fails (the column isn't a direct passthrough
/// from an inner table), this traces through the subquery's SELECT items
/// to find the expression that creates the alias and resolves its base columns.
#[allow(clippy::too_many_arguments)]
fn resolve_subquery_expression_column(
    subquery_alias: &str,
    column: &str,
    stmts: &[Statement],
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> Vec<String> {
    let alias_lower = subquery_alias.to_lowercase();
    let col_lower = column.to_lowercase();

    for stmt in stmts {
        let query = match stmt {
            Statement::Query(q) => q,
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    src
                } else {
                    continue;
                }
            }
            _ => continue,
        };

        if let Some(results) = find_subquery_select_items(
            &query.body,
            &alias_lower,
            &col_lower,
            flat_schema,
            cte_names,
            cte_source_tables,
            subquery_sources,
            dialect,
        ) {
            return results;
        }

        // Also search CTE bodies
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                if let Some(results) = find_subquery_select_items(
                    &cte.query.body,
                    &alias_lower,
                    &col_lower,
                    flat_schema,
                    cte_names,
                    cte_source_tables,
                    subquery_sources,
                    dialect,
                ) {
                    return results;
                }
            }
        }
    }

    vec![]
}

/// Recursively search a SetExpr for a Derived subquery with matching alias,
/// then resolve the column through its SELECT items.
#[allow(clippy::too_many_arguments)]
fn find_subquery_select_items(
    body: &SetExpr,
    alias_lower: &str,
    col_lower: &str,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> Option<Vec<String>> {
    match body {
        SetExpr::Select(sel) => {
            for twj in &sel.from {
                if let Some(r) = search_table_factor_for_subquery(
                    &twj.relation,
                    alias_lower,
                    col_lower,
                    flat_schema,
                    cte_names,
                    cte_source_tables,
                    subquery_sources,
                    dialect,
                ) {
                    return Some(r);
                }
                for join in &twj.joins {
                    if let Some(r) = search_table_factor_for_subquery(
                        &join.relation,
                        alias_lower,
                        col_lower,
                        flat_schema,
                        cte_names,
                        cte_source_tables,
                        subquery_sources,
                        dialect,
                    ) {
                        return Some(r);
                    }
                }
            }
            None
        }
        SetExpr::SetOperation { left, right, .. } => find_subquery_select_items(
            left,
            alias_lower,
            col_lower,
            flat_schema,
            cte_names,
            cte_source_tables,
            subquery_sources,
            dialect,
        )
        .or_else(|| {
            find_subquery_select_items(
                right,
                alias_lower,
                col_lower,
                flat_schema,
                cte_names,
                cte_source_tables,
                subquery_sources,
                dialect,
            )
        }),
        SetExpr::Query(q) => find_subquery_select_items(
            &q.body,
            alias_lower,
            col_lower,
            flat_schema,
            cte_names,
            cte_source_tables,
            subquery_sources,
            dialect,
        ),
        _ => None,
    }
}

#[allow(clippy::too_many_arguments)]
fn search_table_factor_for_subquery(
    factor: &ast::TableFactor,
    alias_lower: &str,
    col_lower: &str,
    flat_schema: &FlatSchema,
    cte_names: &BTreeSet<String>,
    cte_source_tables: &BTreeMap<String, Vec<String>>,
    subquery_sources: &BTreeMap<String, Vec<String>>,
    dialect: SqlDialect,
) -> Option<Vec<String>> {
    match factor {
        ast::TableFactor::Derived {
            alias, subquery, ..
        } => {
            let matches = alias
                .as_ref()
                .is_some_and(|a| a.name.value.to_lowercase() == alias_lower);

            if matches {
                // Found our subquery — resolve the column from its SELECT items
                let items = extract_leftmost_select_items(&subquery.body);
                let mut sub = BTreeMap::new();
                let inner_alias_map = build_alias_map_from_body(&subquery.body, &mut sub);

                for item in &items {
                    let item_name = match item {
                        SelectItem::ExprWithAlias { alias, .. } => alias.value.clone(),
                        SelectItem::UnnamedExpr(Expr::Identifier(ident)) => ident.value.clone(),
                        SelectItem::UnnamedExpr(Expr::CompoundIdentifier(parts)) => {
                            parts.last().map(|p| p.value.clone()).unwrap_or_default()
                        }
                        _ => continue,
                    };

                    if item_name.to_lowercase() != *col_lower {
                        continue;
                    }

                    let mut raw_refs = Vec::new();
                    match item {
                        SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
                            extract_columns_from_expr(e, &mut raw_refs);
                        }
                        _ => {}
                    }

                    let mut resolved = Vec::new();
                    for raw in &raw_refs {
                        // Canonicalize table qualifier before resolving
                        let canon_raw = if let Some(ref tq) = raw.table {
                            if let Some(canon) = find_schema_table(flat_schema, tq) {
                                RawColumnRef {
                                    table: Some(canon.to_string()),
                                    column: raw.column.clone(),
                                }
                            } else {
                                // Try resolving through inner alias map
                                let actual = inner_alias_map
                                    .get(tq.as_str())
                                    .cloned()
                                    .unwrap_or_else(|| tq.clone());
                                if actual != "__subquery__" {
                                    if let Some(canon) = find_schema_table(flat_schema, &actual) {
                                        RawColumnRef {
                                            table: Some(canon.to_string()),
                                            column: raw.column.clone(),
                                        }
                                    } else {
                                        raw.clone()
                                    }
                                } else {
                                    raw.clone()
                                }
                            }
                        } else {
                            raw.clone()
                        };

                        let r = resolve_column_ref(
                            &canon_raw,
                            &inner_alias_map,
                            flat_schema,
                            cte_names,
                            cte_source_tables,
                            subquery_sources,
                            dialect,
                        );
                        resolved.extend(r);
                    }

                    if !resolved.is_empty() {
                        return Some(resolved);
                    }
                }
                return Some(vec![]);
            }

            // Recurse into the subquery body to find nested subqueries
            find_subquery_select_items(
                &subquery.body,
                alias_lower,
                col_lower,
                flat_schema,
                cte_names,
                cte_source_tables,
                subquery_sources,
                dialect,
            )
        }
        ast::TableFactor::NestedJoin {
            table_with_joins, ..
        } => {
            if let Some(r) = search_table_factor_for_subquery(
                &table_with_joins.relation,
                alias_lower,
                col_lower,
                flat_schema,
                cte_names,
                cte_source_tables,
                subquery_sources,
                dialect,
            ) {
                return Some(r);
            }
            for join in &table_with_joins.joins {
                if let Some(r) = search_table_factor_for_subquery(
                    &join.relation,
                    alias_lower,
                    col_lower,
                    flat_schema,
                    cte_names,
                    cte_source_tables,
                    subquery_sources,
                    dialect,
                ) {
                    return Some(r);
                }
            }
            None
        }
        _ => None,
    }
}

// ===========================================================================
// Column alias tracing
//
// When polyglot-sql returns an edge like TABLE.ALIASED_COL where
// ALIASED_COL doesn't exist in TABLE's schema, we trace through the
// query's SELECT lists to find the original column. Also handles
// unresolved subquery/CTE aliases by finding their body in the AST
// and tracing the column through it.
// ===========================================================================

/// Trace a column through a SELECT body to find its base table source.
///
/// Searches SELECT lists for the column as an alias or passthrough,
/// resolves through alias maps, and recurses into subqueries.
/// Returns qualified refs like `"TABLE"."COLUMN"`.
fn trace_column_to_base(
    body: &SetExpr,
    column: &str,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<String> {
    if depth > 10 {
        return vec![];
    }

    match body {
        SetExpr::Select(sel) => trace_column_in_select(sel, column, flat_schema, dialect, depth),
        SetExpr::SetOperation { left, right, .. } => {
            // Try left branch first (column names come from left in SQL),
            // but also check right for expression definitions (UNION ALL
            // branches may define expression columns differently).
            let left_results = trace_column_to_base(left, column, flat_schema, dialect, depth + 1);
            if !left_results.is_empty() {
                return left_results;
            }
            trace_column_to_base(right, column, flat_schema, dialect, depth + 1)
        }
        SetExpr::Query(q) => {
            // Check CTEs first
            if let Some(with) = &q.with {
                for cte in &with.cte_tables {
                    let cte_cols = extract_select_column_names(&cte.query.body);
                    if cte_cols
                        .iter()
                        .any(|c| c.eq_ignore_ascii_case(column) || c == "*")
                    {
                        let results = trace_column_to_base(
                            &cte.query.body,
                            column,
                            flat_schema,
                            dialect,
                            depth + 1,
                        );
                        if !results.is_empty() {
                            return results;
                        }
                    }
                }
            }
            trace_column_to_base(&q.body, column, flat_schema, dialect, depth + 1)
        }
        _ => vec![],
    }
}

/// Trace a column through a specific SELECT statement.
fn trace_column_in_select(
    sel: &ast::Select,
    column: &str,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<String> {
    let mut sub = BTreeMap::new();
    let alias_map = build_alias_map(&sel.from, &mut sub);

    // 1. Check if column exists directly in a schema table from FROM
    for actual_table in alias_map.values() {
        if actual_table == "__subquery__" {
            continue;
        }
        if let Some(canon) = find_schema_table(flat_schema, actual_table)
            && schema_has_column(flat_schema, canon, column)
        {
            return vec![qualify_ref(canon, column, dialect)];
        }
    }

    // 2. Check SELECT list for column as alias
    for item in &sel.projection {
        match item {
            SelectItem::ExprWithAlias { alias, expr } => {
                if alias.value.eq_ignore_ascii_case(column) {
                    let results =
                        trace_expr_to_base_refs(expr, &alias_map, sel, flat_schema, dialect, depth);
                    if !results.is_empty() {
                        return results;
                    }
                }
            }
            SelectItem::UnnamedExpr(expr) => {
                let expr_name = match expr {
                    Expr::Identifier(ident) => Some(ident.value.as_str()),
                    Expr::CompoundIdentifier(parts) => parts.last().map(|p| p.value.as_str()),
                    _ => None,
                };
                if let Some(name) = expr_name
                    && name.eq_ignore_ascii_case(column)
                {
                    // Passthrough or qualified ref — trace it
                    let results =
                        trace_expr_to_base_refs(expr, &alias_map, sel, flat_schema, dialect, depth);
                    if !results.is_empty() {
                        return results;
                    }
                }
            }
            SelectItem::Wildcard(_) => {}
            SelectItem::QualifiedWildcard(obj_name, _) => {
                let qualifier = obj_name.to_string();
                let resolved = alias_map
                    .iter()
                    .find(|(k, _)| k.eq_ignore_ascii_case(&qualifier))
                    .map(|(_, v)| v.as_str());
                if let Some(actual) = resolved
                    && actual != "__subquery__"
                    && let Some(canon) = find_schema_table(flat_schema, actual)
                    && schema_has_column(flat_schema, canon, column)
                {
                    return vec![qualify_ref(canon, column, dialect)];
                }
            }
        }
    }

    // 3. Recurse into derived subqueries as fallback (column might pass through SELECT *)
    {
        for twj in &sel.from {
            if let ast::TableFactor::Derived { subquery, .. } = &twj.relation {
                let results =
                    trace_column_to_base(&subquery.body, column, flat_schema, dialect, depth + 1);
                if !results.is_empty() {
                    return results;
                }
            }
            for join in &twj.joins {
                if let ast::TableFactor::Derived { subquery, .. } = &join.relation {
                    let results = trace_column_to_base(
                        &subquery.body,
                        column,
                        flat_schema,
                        dialect,
                        depth + 1,
                    );
                    if !results.is_empty() {
                        return results;
                    }
                }
            }
        }
    }

    vec![]
}

/// Extract column references from an expression and trace each to base tables.
fn trace_expr_to_base_refs(
    expr: &Expr,
    alias_map: &AliasMap,
    sel: &ast::Select,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<String> {
    let mut refs = Vec::new();
    extract_columns_from_expr(expr, &mut refs);
    let mut results = Vec::new();

    for raw in &refs {
        if let Some(table_qualifier) = &raw.table {
            let resolved = alias_map
                .iter()
                .find(|(k, _)| k.eq_ignore_ascii_case(table_qualifier))
                .map(|(_, v)| v.as_str())
                .unwrap_or(table_qualifier.as_str());

            if resolved == "__subquery__" {
                // Find the subquery body and recurse
                for twj in &sel.from {
                    if let Some(body) = find_subquery_body_in_factor(&twj.relation, table_qualifier)
                    {
                        let traced = trace_column_to_base(
                            body,
                            &raw.column,
                            flat_schema,
                            dialect,
                            depth + 1,
                        );
                        results.extend(traced);
                    }
                    for join in &twj.joins {
                        if let Some(body) =
                            find_subquery_body_in_factor(&join.relation, table_qualifier)
                        {
                            let traced = trace_column_to_base(
                                body,
                                &raw.column,
                                flat_schema,
                                dialect,
                                depth + 1,
                            );
                            results.extend(traced);
                        }
                    }
                }
            } else if let Some(canon) = find_schema_table(flat_schema, resolved)
                && schema_has_column(flat_schema, canon, &raw.column)
            {
                results.push(qualify_ref(canon, &raw.column, dialect));
            }
        } else {
            // Unqualified — search FROM tables
            let mut found = false;
            for actual_table in alias_map.values() {
                if actual_table == "__subquery__" {
                    continue;
                }
                if let Some(canon) = find_schema_table(flat_schema, actual_table)
                    && schema_has_column(flat_schema, canon, &raw.column)
                {
                    results.push(qualify_ref(canon, &raw.column, dialect));
                    found = true;
                    break;
                }
            }
            if !found {
                // Try derived subqueries
                for twj in &sel.from {
                    if let ast::TableFactor::Derived { subquery, .. } = &twj.relation {
                        let traced = trace_column_to_base(
                            &subquery.body,
                            &raw.column,
                            flat_schema,
                            dialect,
                            depth + 1,
                        );
                        if !traced.is_empty() {
                            results.extend(traced);
                            break;
                        }
                    }
                    if found {
                        break;
                    }
                    for join in &twj.joins {
                        if let ast::TableFactor::Derived { subquery, .. } = &join.relation {
                            let traced = trace_column_to_base(
                                &subquery.body,
                                &raw.column,
                                flat_schema,
                                dialect,
                                depth + 1,
                            );
                            if !traced.is_empty() {
                                results.extend(traced);
                                found = true;
                                break;
                            }
                        }
                    }
                }
            }
        }
    }

    results
}

/// Search the entire query AST for a derived table with the given alias,
/// then trace a column through its body to base table sources.
fn find_and_trace_derived(
    body: &SetExpr,
    alias: &str,
    column: &str,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
    depth: usize,
) -> Vec<String> {
    if depth > 10 {
        return vec![];
    }

    match body {
        SetExpr::Select(sel) => {
            for twj in &sel.from {
                // Check if this factor IS the alias we're looking for
                if let Some(sub_body) = find_subquery_body_in_factor(&twj.relation, alias) {
                    return trace_column_to_base(sub_body, column, flat_schema, dialect, 0);
                }
                // Recurse into non-matching derived tables
                if let ast::TableFactor::Derived { subquery, .. } = &twj.relation {
                    let r = find_and_trace_derived(
                        &subquery.body,
                        alias,
                        column,
                        flat_schema,
                        dialect,
                        depth + 1,
                    );
                    if !r.is_empty() {
                        return r;
                    }
                }
                for join in &twj.joins {
                    if let Some(sub_body) = find_subquery_body_in_factor(&join.relation, alias) {
                        return trace_column_to_base(sub_body, column, flat_schema, dialect, 0);
                    }
                    if let ast::TableFactor::Derived { subquery, .. } = &join.relation {
                        let r = find_and_trace_derived(
                            &subquery.body,
                            alias,
                            column,
                            flat_schema,
                            dialect,
                            depth + 1,
                        );
                        if !r.is_empty() {
                            return r;
                        }
                    }
                }
            }
            vec![]
        }
        SetExpr::SetOperation { left, right, .. } => {
            let r = find_and_trace_derived(left, alias, column, flat_schema, dialect, depth + 1);
            if !r.is_empty() {
                return r;
            }
            find_and_trace_derived(right, alias, column, flat_schema, dialect, depth + 1)
        }
        SetExpr::Query(q) => {
            if let Some(with) = &q.with {
                for cte in &with.cte_tables {
                    if cte.alias.name.value.eq_ignore_ascii_case(alias) {
                        return trace_column_to_base(
                            &cte.query.body,
                            column,
                            flat_schema,
                            dialect,
                            0,
                        );
                    }
                }
            }
            find_and_trace_derived(&q.body, alias, column, flat_schema, dialect, depth + 1)
        }
        _ => vec![],
    }
}

/// Resolve a source column that doesn't exist in its table's schema.
/// Searches the AST for where the column is aliased from.
fn trace_aliased_column(
    stmts: &[Statement],
    column: &str,
    _expected_table: &str,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) -> Vec<String> {
    for stmt in stmts {
        let query = match stmt {
            Statement::Query(q) => q,
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    src
                } else {
                    continue;
                }
            }
            _ => continue,
        };
        // Search CTEs
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                let results =
                    trace_column_to_base(&cte.query.body, column, flat_schema, dialect, 0);
                if !results.is_empty() {
                    return results;
                }
            }
        }
        // Search main body
        let results = trace_column_to_base(&query.body, column, flat_schema, dialect, 0);
        if !results.is_empty() {
            return results;
        }
    }
    vec![]
}

/// Resolve a source table that isn't in the schema (subquery/CTE alias).
/// Finds the alias's body in the AST and traces the column through it.
fn trace_through_unknown_source(
    stmts: &[Statement],
    source_alias: &str,
    column: &str,
    flat_schema: &FlatSchema,
    dialect: SqlDialect,
) -> Vec<String> {
    for stmt in stmts {
        let query = match stmt {
            Statement::Query(q) => q,
            Statement::Insert(ins) => {
                if let Some(ref src) = ins.source {
                    src
                } else {
                    continue;
                }
            }
            _ => continue,
        };
        // Check CTEs
        if let Some(with) = &query.with {
            for cte in &with.cte_tables {
                if cte.alias.name.value.eq_ignore_ascii_case(source_alias) {
                    return trace_column_to_base(&cte.query.body, column, flat_schema, dialect, 0);
                }
            }
        }
        // Search for derived table
        let results =
            find_and_trace_derived(&query.body, source_alias, column, flat_schema, dialect, 0);
        if !results.is_empty() {
            return results;
        }
    }
    vec![]
}

// ===========================================================================
// Main entry point
// ===========================================================================

/// Schema-aware column lineage analysis.
///
/// Single entry point for all column lineage modes.  Resolves column
/// references using schema context (CTE resolution, subquery alias
/// resolution, star expansion) and optionally injects shared-condition
/// dependencies (WHERE, JOIN ON, GROUP BY, HAVING).
///
/// Pass `None` for `options` to get projection-only lineage with no
/// default database/schema qualification — the most common case.
///
/// # Examples
///
/// ```ignore
/// // Projection-only (default):
/// column_lineage(sql, SqlDialect::Snowflake, &schema, None)?;
///
/// // With options:
/// let opts = ColumnLineageOptions {
///     default_database: Some("MY_DB".into()),
///     inject_shared: true,
///     ..Default::default()
/// };
/// column_lineage(sql, SqlDialect::Snowflake, &schema, Some(&opts))?;
/// ```
pub fn column_lineage(
    sql: &str,
    dialect: SqlDialect,
    schema: &HashMap<String, serde_json::Value>,
    options: Option<&ColumnLineageOptions>,
) -> Result<CompleteLineageResult, String> {
    let defaults = ColumnLineageOptions::default();
    let opts = options.unwrap_or(&defaults);
    build_column_lineage(
        sql,
        schema,
        dialect,
        opts.inject_shared,
        opts.default_database.as_deref(),
        opts.default_schema.as_deref(),
    )
}

/// AST-based lineage fallback: trace each outer SELECT item through CTEs to
/// base tables. Used when polyglot-sql returns zero edges (deep CTE chains,
/// complex subqueries, etc.).
fn ast_based_lineage_fallback(
    column_dict: &mut HashMap<String, BTreeSet<String>>,
    info: &QueryInfo,
    flat_schema: &FlatSchema,
    cte_dict: &CteDict,
    dialect: SqlDialect,
) {
    for item in &info.select_items {
        match item {
            SelectItem::Wildcard(_) => {
                // Retry star expansion: expand_stars already ran but may have failed
                // because the column_dict was empty (no polyglot edges to key off).
                // Now try expanding directly from the FROM-clause scope.
                expand_star_from_scope(column_dict, info, flat_schema, cte_dict, dialect);

                // If still empty, try a deeper subquery walk:
                // for SELECT * FROM (SELECT * FROM t), we need to recurse into nested subqueries
                if column_dict.is_empty() {
                    expand_star_through_nested_subqueries(
                        column_dict,
                        &info.query.body,
                        info,
                        flat_schema,
                        cte_dict,
                        dialect,
                        0,
                    );
                }
            }
            SelectItem::QualifiedWildcard(name, _) => {
                let prefix = name.to_string();
                let actual = info
                    .alias_map
                    .iter()
                    .find(|(k, _)| k.eq_ignore_ascii_case(&prefix))
                    .map(|(_, v)| v.clone())
                    .unwrap_or_else(|| prefix.clone());
                let is_cte = info
                    .cte_names
                    .iter()
                    .any(|c| c.eq_ignore_ascii_case(&actual));
                if is_cte {
                    expand_cte_star(&actual, column_dict, info, flat_schema, cte_dict, dialect);
                } else {
                    expand_table_star(&actual, column_dict, info, flat_schema, dialect);
                }
            }
            SelectItem::UnnamedExpr(expr) | SelectItem::ExprWithAlias { expr, .. } => {
                // Get output column name
                let col_name = match item {
                    SelectItem::ExprWithAlias { alias, .. } => {
                        normalize_column_name(&alias.value, dialect)
                    }
                    SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
                        normalize_column_name(&ident.value, dialect)
                    }
                    SelectItem::UnnamedExpr(Expr::CompoundIdentifier(parts)) => {
                        normalize_column_name(
                            &parts.last().map(|p| p.value.clone()).unwrap_or_default(),
                            dialect,
                        )
                    }
                    _ => continue,
                };

                // Extract column refs from the expression
                let mut raw_refs = Vec::new();
                extract_columns_from_expr(expr, &mut raw_refs);

                if raw_refs.is_empty() {
                    continue;
                }

                // Resolve using CTE dict + schema (Python's _find_alias_col approach)
                for raw in &raw_refs {
                    let resolved = resolve_ref_with_cte_dict(
                        raw,
                        &info.alias_map,
                        flat_schema,
                        cte_dict,
                        dialect,
                        None,
                    );
                    for r in &resolved {
                        column_dict
                            .entry(col_name.clone())
                            .or_default()
                            .insert(r.clone());
                    }
                }

                // If direct resolution failed, try CTE expression tracing
                if !column_dict.contains_key(&col_name) || column_dict[&col_name].is_empty() {
                    // Check if any FROM-clause table is a CTE
                    for actual_table in info.alias_map.values() {
                        let is_cte = info
                            .cte_names
                            .iter()
                            .any(|c| c.eq_ignore_ascii_case(actual_table));
                        if is_cte {
                            // Try resolving through CTE column
                            let cte_resolved = resolve_cte_column(
                                actual_table,
                                &col_name,
                                &info.cte_source_tables,
                                flat_schema,
                                dialect,
                            );
                            if !cte_resolved.is_empty() {
                                for r in &cte_resolved {
                                    column_dict
                                        .entry(col_name.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                                break;
                            }
                            // Try CTE expression column
                            let expr_deps = resolve_cte_expression_column(
                                actual_table,
                                &col_name,
                                &info.query,
                                flat_schema,
                                &info.cte_names,
                                &info.cte_source_tables,
                                &info.subquery_source_tables,
                                dialect,
                            );
                            if all_refs_in_schema(&expr_deps, flat_schema) {
                                for dep in &expr_deps {
                                    column_dict
                                        .entry(col_name.clone())
                                        .or_default()
                                        .insert(dep.clone());
                                }
                                break;
                            }
                        }
                    }
                }
            }
        }
    }
}

/// Core lineage builder shared by both projection and complete modes.
///
/// When `inject_shared` is false, skips shared-condition injection (projection-only).
/// When `inject_shared` is true, injects WHERE/JOIN/GROUP BY/HAVING dependencies.
///
/// `default_database` / `default_schema`: when provided, unqualified table
/// references in source refs are qualified with these defaults. Passed through
/// to `CompleteLineageResult.default_database` / `default_schema`.
fn build_column_lineage(
    sql: &str,
    schema: &HashMap<String, serde_json::Value>,
    dialect: SqlDialect,
    inject_shared: bool,
    default_database: Option<&str>,
    default_schema: Option<&str>,
) -> Result<CompleteLineageResult, String> {
    let clean_sql = strip_comments(sql);

    // Step 1: Get projection edges from polyglot-sql.
    // Always uses values_only mode: excludes CASE WHEN conditions,
    // IFF condition args, and window PARTITION BY / ORDER BY from edges.
    // This matches Snowflake ACCESS_USAGE "direct" lineage semantics.
    let base = polyglot_column_lineage(&clean_sql, dialect)?;

    // Step 2: Parse SQL and analyze query structure
    // Try dialect-specific parser first; fall back to BigQuery parser (most permissive,
    // handles trailing commas in SELECT lists that other dialects reject).
    let parser_dialect = get_parser_dialect(dialect);
    let stmts = Parser::parse_sql(parser_dialect.as_ref(), &clean_sql)
        .or_else(|_| {
            let fallback = get_parser_dialect(SqlDialect::BigQuery);
            Parser::parse_sql(fallback.as_ref(), &clean_sql)
        })
        .ok();
    let flat_schema = flatten_schema(schema);
    let info = stmts.as_ref().and_then(|s| analyze_query(s));

    // Step 2b: Build CTE + subquery dictionary — pre-compute column-to-source mappings
    let cte_dict = if let Some(ref info) = info {
        let mut dict = build_cte_dict(&info.query, &flat_schema, dialect);
        extend_with_subquery_dict(&info.query, &flat_schema, &mut dict, dialect);
        dict
    } else {
        CteDict::new()
    };

    // Step 3: Build column_dict from polyglot edges
    let mut column_dict: HashMap<String, BTreeSet<String>> = HashMap::new();
    let mut source_tables: BTreeSet<String> = BTreeSet::new();

    let cte_names = info
        .as_ref()
        .map(|i| &i.cte_names)
        .cloned()
        .unwrap_or_default();
    let cte_source_tables = info
        .as_ref()
        .map(|i| &i.cte_source_tables)
        .cloned()
        .unwrap_or_default();
    let subquery_sources = info
        .as_ref()
        .map(|i| &i.subquery_source_tables)
        .cloned()
        .unwrap_or_default();
    let is_set_op = info.as_ref().is_some_and(|i| i.is_set_operation);
    let left_tables = info
        .as_ref()
        .map(|i| &i.left_side_tables)
        .cloned()
        .unwrap_or_default();
    let flatten_inputs = info
        .as_ref()
        .map(|i| &i.lateral_flatten_inputs)
        .cloned()
        .unwrap_or_default();

    for edge in &base.edges {
        let source_table = &edge.source_table;
        let source_column = &edge.source_column;
        let target = normalize_column_name(&edge.target_column, dialect);

        // Resolve LATERAL FLATTEN aliases (e.g., F.VALUE from
        // `LATERAL FLATTEN(input => col) F`) to the flattened input column.
        if !flatten_inputs.is_empty() {
            let source_bare = source_table.split('.').next_back().unwrap_or(source_table);
            if let Some(input_refs) = flatten_inputs
                .iter()
                .find(|(a, _)| {
                    a.eq_ignore_ascii_case(source_table) || a.eq_ignore_ascii_case(source_bare)
                })
                .map(|(_, refs)| refs.clone())
            {
                // Resolve the FLATTEN input column refs to base table sources
                for raw in &input_refs {
                    let resolved = resolve_ref_with_cte_dict(
                        raw,
                        &info
                            .as_ref()
                            .map(|i| &i.alias_map)
                            .cloned()
                            .unwrap_or_default(),
                        &flat_schema,
                        &cte_dict,
                        dialect,
                        None,
                    );
                    for r in &resolved {
                        column_dict
                            .entry(target.clone())
                            .or_default()
                            .insert(r.clone());
                    }
                }
                continue;
            }
        }

        match edge.terminal_kind {
            TerminalKind::Table => {
                // 1. CTE resolution (case-insensitive)
                if let Some(cte_name) = find_cte_name(&cte_names, source_table) {
                    let cte_name = cte_name.clone();
                    // For set operations: only keep CTEs from the left side
                    if is_set_op && !left_tables.is_empty() {
                        let is_left = left_tables
                            .iter()
                            .any(|t| t.eq_ignore_ascii_case(&cte_name));
                        if !is_left {
                            continue;
                        }
                    }

                    // Try CTE dict first — pre-computed column→base table mappings
                    if let Some(cte_cols) = find_cte_in_dict(&cte_dict, &cte_name) {
                        let col_upper = source_column.to_uppercase();
                        if let Some(sources) = cte_cols.get(&col_upper).or_else(|| {
                            cte_cols
                                .iter()
                                .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                .map(|(_, v)| v)
                        }) {
                            if !sources.is_empty() {
                                for r in sources {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                                continue;
                            }
                            // Column in CTE dict but no base table sources.
                            // Skip to avoid emitting CTE name as source.
                            continue;
                        }
                    }

                    // Fallback to existing CTE resolution
                    let resolved = resolve_cte_column(
                        &cte_name,
                        source_column,
                        &cte_source_tables,
                        &flat_schema,
                        dialect,
                    );
                    if resolved.is_empty() {
                        // Column not directly mapped — trace expression alias
                        if let Some(ref info) = info {
                            let expr_deps = resolve_cte_expression_column(
                                &cte_name,
                                source_column,
                                &info.query,
                                &flat_schema,
                                &cte_names,
                                &cte_source_tables,
                                &subquery_sources,
                                dialect,
                            );
                            if all_refs_in_schema(&expr_deps, &flat_schema) {
                                for dep in &expr_deps {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(dep.clone());
                                }
                                continue;
                            }
                        }
                        // CTE column untraceable — skip rather than emit CTE name
                        continue;
                    } else {
                        for r in &resolved {
                            column_dict
                                .entry(target.clone())
                                .or_default()
                                .insert(r.clone());
                        }
                    }
                    continue;
                }

                // 2. Subquery alias resolution — try dict first, then schema-based
                if let Some(sub_alias) = find_subquery_source(&subquery_sources, source_table) {
                    let sub_alias = sub_alias.clone();

                    // Try subquery dict first (pre-computed column→sources)
                    if let Some(sq_cols) = find_cte_in_dict(&cte_dict, &sub_alias) {
                        let col_upper = source_column.to_uppercase();
                        if let Some(sources) = sq_cols.get(&col_upper).or_else(|| {
                            sq_cols
                                .iter()
                                .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                .map(|(_, v)| v)
                        }) {
                            if !sources.is_empty() {
                                for r in sources {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                                continue;
                            }
                            // Column in subquery dict but no base table sources.
                            // Skip to avoid emitting subquery alias as source.
                            continue;
                        }
                    }

                    let resolved = resolve_subquery_column(
                        &sub_alias,
                        source_column,
                        &subquery_sources,
                        &flat_schema,
                        dialect,
                    );
                    if !resolved.is_empty() {
                        for r in &resolved {
                            column_dict
                                .entry(target.clone())
                                .or_default()
                                .insert(r.clone());
                        }
                        continue;
                    }
                    // Column not in inner tables — trace expression alias
                    if let Some(ref s) = stmts {
                        let expr_deps = resolve_subquery_expression_column(
                            &sub_alias,
                            source_column,
                            s,
                            &flat_schema,
                            &cte_names,
                            &cte_source_tables,
                            &subquery_sources,
                            dialect,
                        );
                        if all_refs_in_schema(&expr_deps, &flat_schema) {
                            for dep in &expr_deps {
                                column_dict
                                    .entry(target.clone())
                                    .or_default()
                                    .insert(dep.clone());
                            }
                            continue;
                        }
                    }
                }

                // 3. Alias map lookup (case-insensitive)
                if let Some(ref info) = info
                    && let Some(actual) = find_alias(&info.alias_map, source_table)
                {
                    if actual == "__subquery__" {
                        // Try dict first for subquery alias
                        if let Some(sq_cols) = find_cte_in_dict(&cte_dict, source_table) {
                            let col_upper = source_column.to_uppercase();
                            if let Some(sources) = sq_cols.get(&col_upper).or_else(|| {
                                sq_cols
                                    .iter()
                                    .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                    .map(|(_, v)| v)
                            }) && !sources.is_empty()
                            {
                                for r in sources {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                                continue;
                            }
                        }
                        let resolved = resolve_subquery_column(
                            source_table,
                            source_column,
                            &subquery_sources,
                            &flat_schema,
                            dialect,
                        );
                        if !resolved.is_empty() {
                            for r in &resolved {
                                column_dict
                                    .entry(target.clone())
                                    .or_default()
                                    .insert(r.clone());
                            }
                            continue;
                        }
                        // Trace expression alias through subquery body
                        if let Some(ref s) = stmts {
                            let expr_deps = resolve_subquery_expression_column(
                                source_table,
                                source_column,
                                s,
                                &flat_schema,
                                &cte_names,
                                &cte_source_tables,
                                &subquery_sources,
                                dialect,
                            );
                            if all_refs_in_schema(&expr_deps, &flat_schema) {
                                for dep in &expr_deps {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(dep.clone());
                                }
                                continue;
                            }
                        }
                    } else if let Some(cte_name) = find_cte_name(&cte_names, actual) {
                        // Alias points to a CTE — try CTE dict first
                        let cte_name = cte_name.clone();
                        let mut found = false;
                        if let Some(cte_cols) = find_cte_in_dict(&cte_dict, &cte_name) {
                            let col_upper = source_column.to_uppercase();
                            if let Some(sources) = cte_cols.get(&col_upper).or_else(|| {
                                cte_cols
                                    .iter()
                                    .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                    .map(|(_, v)| v)
                            }) && !sources.is_empty()
                            {
                                for r in sources {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                                found = true;
                            }
                        }
                        if found {
                            continue;
                        }
                        let resolved = resolve_cte_column(
                            &cte_name,
                            source_column,
                            &cte_source_tables,
                            &flat_schema,
                            dialect,
                        );
                        if !resolved.is_empty() {
                            for r in &resolved {
                                column_dict
                                    .entry(target.clone())
                                    .or_default()
                                    .insert(r.clone());
                            }
                            continue;
                        }
                    } else if let Some(canon) = find_schema_table(&flat_schema, actual) {
                        // Alias points to a real table
                        let source_ref = qualify_ref(canon, source_column, dialect);
                        column_dict.entry(target).or_default().insert(source_ref);
                        continue;
                    }
                }

                // 3.5. CTE-internal alias resolution
                // Polyglot-sql may output CTE-internal aliases (e.g., `X` from
                // `FROM global_audience X` inside a CTE). Check all CTE alias
                // maps to resolve these to actual table names.
                if let Some(ref info) = info {
                    let mut resolved_via_cte_alias = false;
                    for cte_am in info.cte_alias_maps.values() {
                        if let Some(actual) = find_alias(cte_am, source_table) {
                            if actual == "__subquery__" {
                                // Try dict for inner subquery
                                if let Some(sq_cols) = find_cte_in_dict(&cte_dict, source_table) {
                                    let col_upper = source_column.to_uppercase();
                                    if let Some(sources) = sq_cols.get(&col_upper).or_else(|| {
                                        sq_cols
                                            .iter()
                                            .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                            .map(|(_, v)| v)
                                    }) && !sources.is_empty()
                                    {
                                        for r in sources {
                                            column_dict
                                                .entry(target.clone())
                                                .or_default()
                                                .insert(r.clone());
                                        }
                                        resolved_via_cte_alias = true;
                                        break;
                                    }
                                }
                            } else if let Some(canon) = find_schema_table(&flat_schema, actual) {
                                let source_ref = qualify_ref(canon, source_column, dialect);
                                column_dict
                                    .entry(target.clone())
                                    .or_default()
                                    .insert(source_ref);
                                resolved_via_cte_alias = true;
                                break;
                            } else if let Some(cte_n) = find_cte_name(&cte_names, actual) {
                                // Alias points to another CTE — resolve via dict
                                let cte_n = cte_n.clone();
                                if let Some(cte_cols) = find_cte_in_dict(&cte_dict, &cte_n) {
                                    let col_upper = source_column.to_uppercase();
                                    if let Some(sources) = cte_cols.get(&col_upper).or_else(|| {
                                        cte_cols
                                            .iter()
                                            .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                            .map(|(_, v)| v)
                                    }) && !sources.is_empty()
                                    {
                                        for r in sources {
                                            column_dict
                                                .entry(target.clone())
                                                .or_default()
                                                .insert(r.clone());
                                        }
                                        resolved_via_cte_alias = true;
                                        break;
                                    }
                                }
                            }
                        }
                    }
                    if resolved_via_cte_alias {
                        continue;
                    }
                }

                // 4. Regular table: set-op left-side filtering
                if is_set_op && !left_tables.is_empty() {
                    let is_left = left_tables
                        .iter()
                        .any(|t| t.eq_ignore_ascii_case(source_table));
                    if !is_left {
                        continue;
                    }
                }

                // 5. Check if source_table exists in schema
                if let Some(canon) = find_schema_table(&flat_schema, source_table) {
                    source_tables.insert(canon.to_string());
                    if schema_has_column(&flat_schema, canon, source_column) {
                        // Column exists in table — use directly
                        let source_ref = qualify_ref(canon, source_column, dialect);
                        column_dict.entry(target).or_default().insert(source_ref);
                    } else {
                        // Column doesn't exist in table — trace alias through AST
                        let traced = if let Some(ref s) = stmts {
                            trace_aliased_column(s, source_column, canon, &flat_schema, dialect)
                        } else {
                            vec![]
                        };
                        if !traced.is_empty() {
                            for r in &traced {
                                column_dict
                                    .entry(target.clone())
                                    .or_default()
                                    .insert(r.clone());
                            }
                        } else {
                            // Couldn't trace — use as-is
                            let source_ref = qualify_ref(canon, source_column, dialect);
                            column_dict.entry(target).or_default().insert(source_ref);
                        }
                    }
                } else {
                    // 5b. Source table not in schema — try dict (subquery alias)
                    if let Some(sq_cols) = find_cte_in_dict(&cte_dict, source_table) {
                        let col_upper = source_column.to_uppercase();
                        if let Some(sources) = sq_cols.get(&col_upper).or_else(|| {
                            sq_cols
                                .iter()
                                .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                .map(|(_, v)| v)
                        }) {
                            if !sources.is_empty() {
                                for r in sources {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                                continue;
                            }
                            // Column in dict but empty sources — skip
                            continue;
                        }
                    }

                    // 6. Source table not in schema — try AST tracing first
                    let traced = if let Some(ref s) = stmts {
                        trace_through_unknown_source(
                            s,
                            source_table,
                            source_column,
                            &flat_schema,
                            dialect,
                        )
                    } else {
                        vec![]
                    };
                    if !traced.is_empty() {
                        for r in &traced {
                            column_dict
                                .entry(target.clone())
                                .or_default()
                                .insert(r.clone());
                        }
                    } else {
                        // Fallback: column search
                        let matching = find_tables_with_column(&flat_schema, source_column);
                        if matching.len() == 1 {
                            let source_ref = qualify_ref(matching[0], source_column, dialect);
                            column_dict.entry(target).or_default().insert(source_ref);
                        } else {
                            let source_ref = edge_to_qualified_ref(edge, dialect);
                            column_dict.entry(target).or_default().insert(source_ref);
                        }
                    }
                }
            }
            TerminalKind::Star | TerminalKind::Placeholder => {
                if is_set_op && !left_tables.is_empty() {
                    let table_lower = source_table.to_lowercase();
                    let is_left = left_tables.iter().any(|t| t.to_lowercase() == table_lower);
                    if !is_left {
                        continue;
                    }
                }
                let source_ref = edge_to_qualified_ref(edge, dialect);
                column_dict.entry(target).or_default().insert(source_ref);
            }
            TerminalKind::Unknown => {
                if source_table.is_empty() {
                    // Empty source_table — resolve from schema
                    let matching = find_tables_with_column(&flat_schema, source_column);
                    if let Some(&table) = matching.first() {
                        let source_ref = qualify_ref(table, source_column, dialect);
                        column_dict
                            .entry(target.clone())
                            .or_default()
                            .insert(source_ref);
                    }
                } else {
                    // Non-empty source_table with Unknown kind
                    let mut resolved = false;

                    // a) CTE resolution — try CTE dict first, then fallback
                    if let Some(cte_name) = find_cte_name(&cte_names, source_table) {
                        let cte_name = cte_name.clone();

                        // Try CTE dict first
                        if let Some(cte_cols) = find_cte_in_dict(&cte_dict, &cte_name) {
                            let col_upper = source_column.to_uppercase();
                            if let Some(sources) = cte_cols.get(&col_upper).or_else(|| {
                                cte_cols
                                    .iter()
                                    .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                    .map(|(_, v)| v)
                            }) {
                                if !sources.is_empty() {
                                    for dep in sources {
                                        column_dict
                                            .entry(target.clone())
                                            .or_default()
                                            .insert(dep.clone());
                                    }
                                    resolved = true;
                                } else {
                                    // Column in CTE dict but no base table sources.
                                    // Mark resolved to avoid fallback emitting CTE name.
                                    resolved = true;
                                }
                            }
                        }

                        if !resolved {
                            let r = resolve_cte_column(
                                &cte_name,
                                source_column,
                                &cte_source_tables,
                                &flat_schema,
                                dialect,
                            );
                            if !r.is_empty() {
                                for dep in &r {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(dep.clone());
                                }
                                resolved = true;
                            }
                        }
                    }

                    // a2) Nested CTE resolution — CTEs defined inside derived
                    // subqueries are in the CTE dict but not in cte_names.
                    // Check the dict directly for any source not yet resolved.
                    if !resolved && let Some(cte_cols) = find_cte_in_dict(&cte_dict, source_table) {
                        let col_upper = source_column.to_uppercase();
                        if let Some(sources) = cte_cols.get(&col_upper).or_else(|| {
                            cte_cols
                                .iter()
                                .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                .map(|(_, v)| v)
                        }) {
                            if !sources.is_empty() {
                                for dep in sources {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(dep.clone());
                                }
                                resolved = true;
                            } else {
                                resolved = true;
                            }
                        }
                    }

                    // b) Subquery alias resolution — try dict first, then fallback
                    if !resolved
                        && let Some(sub_alias) =
                            find_subquery_source(&subquery_sources, source_table)
                    {
                        let sub_alias = sub_alias.clone();

                        // Try subquery dict first
                        if let Some(sq_cols) = find_cte_in_dict(&cte_dict, &sub_alias) {
                            let col_upper = source_column.to_uppercase();
                            if let Some(sources) = sq_cols.get(&col_upper).or_else(|| {
                                sq_cols
                                    .iter()
                                    .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                    .map(|(_, v)| v)
                            }) {
                                if !sources.is_empty() {
                                    for dep in sources {
                                        column_dict
                                            .entry(target.clone())
                                            .or_default()
                                            .insert(dep.clone());
                                    }
                                    resolved = true;
                                } else {
                                    // Column exists in dict but has empty sources
                                    // (untraceable to base table). Mark resolved to
                                    // prevent fallback from emitting intermediate names.
                                    resolved = true;
                                }
                            }
                        }

                        if !resolved {
                            let r = resolve_subquery_column(
                                &sub_alias,
                                source_column,
                                &subquery_sources,
                                &flat_schema,
                                dialect,
                            );
                            if !r.is_empty() {
                                for dep in &r {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(dep.clone());
                                }
                                resolved = true;
                            } else if let Some(ref s) = stmts {
                                let expr_deps = resolve_subquery_expression_column(
                                    &sub_alias,
                                    source_column,
                                    s,
                                    &flat_schema,
                                    &cte_names,
                                    &cte_source_tables,
                                    &subquery_sources,
                                    dialect,
                                );
                                if all_refs_in_schema(&expr_deps, &flat_schema) {
                                    for dep in &expr_deps {
                                        column_dict
                                            .entry(target.clone())
                                            .or_default()
                                            .insert(dep.clone());
                                    }
                                    resolved = true;
                                }
                            }
                        } // if !resolved (inner fallback)
                    }

                    // c) Alias map lookup (case-insensitive)
                    if !resolved
                        && let Some(ref info) = info
                        && let Some(actual) = find_alias(&info.alias_map, source_table)
                    {
                        if actual == "__subquery__" {
                            // Try dict first for subquery alias
                            if let Some(sq_cols) = find_cte_in_dict(&cte_dict, source_table) {
                                let col_upper = source_column.to_uppercase();
                                if let Some(sources) = sq_cols.get(&col_upper).or_else(|| {
                                    sq_cols
                                        .iter()
                                        .find(|(k, _)| k.eq_ignore_ascii_case(source_column))
                                        .map(|(_, v)| v)
                                }) {
                                    if !sources.is_empty() {
                                        for dep in sources {
                                            column_dict
                                                .entry(target.clone())
                                                .or_default()
                                                .insert(dep.clone());
                                        }
                                        resolved = true;
                                    } else {
                                        // Column in dict but empty — skip
                                        resolved = true;
                                    }
                                }
                            }
                            if !resolved {
                                let r = resolve_subquery_column(
                                    source_table,
                                    source_column,
                                    &subquery_sources,
                                    &flat_schema,
                                    dialect,
                                );
                                if !r.is_empty() {
                                    for dep in &r {
                                        column_dict
                                            .entry(target.clone())
                                            .or_default()
                                            .insert(dep.clone());
                                    }
                                    resolved = true;
                                }
                            }
                        } else if let Some(canon) = find_schema_table(&flat_schema, actual) {
                            let source_ref = qualify_ref(canon, source_column, dialect);
                            column_dict
                                .entry(target.clone())
                                .or_default()
                                .insert(source_ref);
                            resolved = true;
                        }
                    }

                    // d) Schema lookup / column search fallback
                    if !resolved {
                        if let Some(canon) = find_schema_table(&flat_schema, source_table) {
                            if schema_has_column(&flat_schema, canon, source_column) {
                                let source_ref = qualify_ref(canon, source_column, dialect);
                                column_dict
                                    .entry(target.clone())
                                    .or_default()
                                    .insert(source_ref);
                            } else {
                                // Column aliased — trace through AST
                                let traced = if let Some(ref s) = stmts {
                                    trace_aliased_column(
                                        s,
                                        source_column,
                                        canon,
                                        &flat_schema,
                                        dialect,
                                    )
                                } else {
                                    vec![]
                                };
                                if !traced.is_empty() {
                                    for r in &traced {
                                        column_dict
                                            .entry(target.clone())
                                            .or_default()
                                            .insert(r.clone());
                                    }
                                } else {
                                    let source_ref = qualify_ref(canon, source_column, dialect);
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(source_ref);
                                }
                            }
                        } else {
                            // Source table not in schema — try AST tracing
                            let traced = if let Some(ref s) = stmts {
                                trace_through_unknown_source(
                                    s,
                                    source_table,
                                    source_column,
                                    &flat_schema,
                                    dialect,
                                )
                            } else {
                                vec![]
                            };
                            if !traced.is_empty() {
                                for r in &traced {
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(r.clone());
                                }
                            } else {
                                let matching = find_tables_with_column(&flat_schema, source_column);
                                if matching.len() == 1 {
                                    let source_ref =
                                        qualify_ref(matching[0], source_column, dialect);
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(source_ref);
                                } else {
                                    let source_ref = edge_to_qualified_ref(edge, dialect);
                                    column_dict
                                        .entry(target.clone())
                                        .or_default()
                                        .insert(source_ref);
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    // Step 4: Expand star columns and handle aggregates
    if let Some(ref info) = info {
        expand_stars(&mut column_dict, info, &flat_schema, &cte_dict, dialect);
        supplement_missing_columns(&mut column_dict, info, &flat_schema, &cte_dict, dialect);

        // Step 4a: AST-based fallback when polyglot returned zero useful edges.
        // Must run BEFORE UNION supplementation so column_dict has entries to merge into.
        if column_dict.is_empty() {
            ast_based_lineage_fallback(&mut column_dict, info, &flat_schema, &cte_dict, dialect);
        }

        // Step 4.5: UNION branch supplementation.
        // Polyglot only generates edges from the leftmost UNION branch. For each
        // non-leftmost branch, resolve SELECT item expressions and merge sources.
        if info.is_set_operation {
            supplement_union_branches(
                &mut column_dict,
                &info.query.body,
                &info.select_items,
                &flat_schema,
                &cte_dict,
                dialect,
            );
        }

        // Step 4.6: FROM-subquery UNION supplementation.
        // Polyglot doesn't trace aggregate-wrapped columns through UNION branches
        // in FROM subqueries (e.g., SELECT * FROM (SELECT MAX(x) UNION ALL SELECT MAX(y))).
        // Walk all FROM subqueries recursively to find UNION bodies and merge sources.
        supplement_from_subquery_unions(
            &mut column_dict,
            &info.query.body,
            &flat_schema,
            &cte_dict,
            dialect,
        );

        // Step 4b: INSERT INTO column remapping
        // For INSERT INTO queries, remap SELECT output columns to INSERT target columns
        // positionally, and discard any extra SELECT columns that exceed the target list.
        if !info.insert_target_columns.is_empty() {
            remap_insert_columns(&mut column_dict, info, &flat_schema, dialect);
        } else if let Some(ref s) = stmts {
            // INSERT without explicit column list: infer from schema
            let inferred = infer_insert_target_columns(s, &flat_schema, dialect);
            if !inferred.is_empty() && inferred.len() < column_dict.len() {
                // Only filter if we have fewer target cols than actual cols
                // (otherwise we might be losing data)
                let keep: BTreeSet<String> = inferred.into_iter().collect();
                let extras: Vec<String> = column_dict
                    .keys()
                    .filter(|k| !keep.contains(k.as_str()))
                    .cloned()
                    .collect();
                for k in extras {
                    column_dict.remove(&k);
                }
            }
        }
    }

    // Step 5-6: Extract and inject shared conditions (complete mode only)
    //
    // shared_deps lives outside the if-block so we can classify direct vs
    // indirect when building the flat lineage list later.
    let mut shared_deps = BTreeSet::new();
    // Snapshot direct sources before injection — only when we actually inject.
    // Avoids the clone in projection mode (inject_shared=false).
    let mut direct_snapshot: Option<HashMap<String, BTreeSet<String>>> = if inject_shared {
        if let Some(ref info) = info {
            let ctx = ResolutionCtx {
                flat_schema: &flat_schema,
                cte_names: &cte_names,
                cte_source_tables: &cte_source_tables,
                subquery_sources: &subquery_sources,
                dialect,
            };

            shared_deps = extract_and_resolve_shared_deps(
                &info.query,
                &info.used_ctes,
                &subquery_sources,
                &info.cte_alias_maps,
                &ctx,
            );

            // GROUP BY positional references
            let group_by_deps = resolve_group_by_positional(
                &info.query,
                &flat_schema,
                &cte_names,
                &cte_source_tables,
                &subquery_sources,
                dialect,
            );
            shared_deps.extend(group_by_deps);
        }

        // Snapshot BEFORE injection so we can tell direct from indirect later
        let snap = column_dict.clone();

        if !shared_deps.is_empty() {
            for sources in column_dict.values_mut() {
                for dep in &shared_deps {
                    sources.insert(dep.clone());
                }
            }
        }

        Some(snap)
    } else {
        None
    };

    // Post-filter: strip window PARTITION BY / ORDER BY columns (projection
    // mode only). CASE/IFF conditions are included in both All and ValuesOnly
    // (they're direct lineage sources), so the only difference is window functions.
    if let Some(ref info) = info {
        let mut target_condition_cols: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
        for (i, item) in info.select_items.iter().enumerate() {
            let target_name = match select_item_name(item, i, dialect) {
                Some(n) => n,
                None => continue,
            };
            let expr = match item {
                SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => e,
                _ => continue,
            };

            let mut all_refs = Vec::new();
            extract_columns_from_expr(expr, &mut all_refs);
            let mut value_refs = Vec::new();
            extract_columns_from_expr_values_only(expr, &mut value_refs);

            // Resolve CASE/IFF condition refs and add to column_dict — polyglot
            // may not produce edges for condition columns that our AST extraction
            // correctly identifies as direct lineage sources.
            if value_refs.is_empty() {
                // Truly all-literal expression: fallback to All mode
                let mut resolved = BTreeSet::new();
                for raw in &all_refs {
                    let refs = resolve_ref_with_cte_dict(
                        raw,
                        &info.alias_map,
                        &flat_schema,
                        &cte_dict,
                        dialect,
                        None,
                    );
                    resolved.extend(refs);
                }
                if !resolved.is_empty() {
                    column_dict.entry(target_name).or_default().extend(resolved);
                }
                continue;
            }

            let all_cols: BTreeSet<String> = all_refs
                .iter()
                .map(|r| normalize_column_name(&r.column, dialect))
                .collect();
            let mut value_cols: BTreeSet<String> = value_refs
                .iter()
                .map(|r| normalize_column_name(&r.column, dialect))
                .collect();

            // In complete mode, preserve PARTITION BY columns as value-producing
            // dependencies rather than stripping them as condition-only.
            if inject_shared {
                let mut partition_refs = Vec::new();
                extract_window_partition_refs(expr, &mut partition_refs);
                if !partition_refs.is_empty() {
                    let mut resolved = BTreeSet::new();
                    for raw in &partition_refs {
                        let r = resolve_ref_with_cte_dict(
                            raw,
                            &info.alias_map,
                            &flat_schema,
                            &cte_dict,
                            dialect,
                            None,
                        );
                        resolved.extend(r);
                    }
                    if !resolved.is_empty() {
                        column_dict
                            .entry(target_name.clone())
                            .or_default()
                            .extend(resolved);
                    }
                    for r in &partition_refs {
                        value_cols.insert(normalize_column_name(&r.column, dialect));
                    }
                }
            }

            let condition_only: BTreeSet<String> =
                all_cols.difference(&value_cols).cloned().collect();

            if !condition_only.is_empty() {
                target_condition_cols.insert(target_name, condition_only);
            }
        }

        // Strip condition-only source refs from column_dict (now only window
        // PARTITION BY / ORDER BY columns in projection mode).
        for (target, sources) in column_dict.iter_mut() {
            if let Some(cond_cols) = target_condition_cols.get(target) {
                let to_remove: Vec<String> = sources
                    .iter()
                    .filter(|ref_str| {
                        let parts: Vec<&str> = ref_str.split('.').collect();
                        if parts.len() >= 2 {
                            let col = parts.last().unwrap().trim_matches('"');
                            cond_cols.contains(&normalize_column_name(col, dialect))
                        } else {
                            false
                        }
                    })
                    .cloned()
                    .collect();
                if !to_remove.is_empty() && to_remove.len() < sources.len() {
                    for s in &to_remove {
                        sources.remove(s);
                    }
                }
            }
        }
    }

    // Filter unresolved CTE/subquery alias source references.
    // When a source table is a known CTE or subquery alias that we should have
    // resolved to base tables but didn't, strip it — the raw alias ref is noise.
    // Only strip if other valid (in-schema) sources exist for the same column.
    if let Some(ref info) = info {
        for sources in column_dict.values_mut() {
            let unresolved: Vec<String> = sources
                .iter()
                .filter(|ref_str| {
                    let parts: Vec<&str> = ref_str.split('.').collect();
                    if parts.len() < 2 {
                        return false;
                    }
                    let table_part = parts[..parts.len() - 1]
                        .iter()
                        .map(|p| p.trim_matches('"'))
                        .collect::<Vec<_>>()
                        .join(".");
                    // Only strip if NOT in schema AND is a known CTE or subquery alias
                    if find_schema_table(&flat_schema, &table_part).is_some() {
                        return false;
                    }
                    let is_cte = info
                        .cte_names
                        .iter()
                        .any(|c| c.eq_ignore_ascii_case(&table_part));
                    let is_subquery = info
                        .subquery_source_tables
                        .keys()
                        .any(|k| k.eq_ignore_ascii_case(&table_part));
                    let is_alias = info
                        .alias_map
                        .keys()
                        .any(|k| k.eq_ignore_ascii_case(&table_part));
                    is_cte || is_subquery || is_alias
                })
                .cloned()
                .collect();
            // Only remove if doing so wouldn't empty the source set
            if !unresolved.is_empty() && unresolved.len() < sources.len() {
                for s in &unresolved {
                    sources.remove(s);
                }
            }
        }
    }
    // CTE dict post-narrowing: when CTE dict provides more precise source info
    // than polyglot edges, remove same-table extras. Only apply in projection
    // mode — complete mode injects shared conditions that CTE dict lacks.
    if !inject_shared {
        for (target_col, sources) in column_dict.iter_mut() {
            let col_upper = target_col.to_uppercase();
            let mut cte_source_sets: Vec<&BTreeSet<String>> = Vec::new();

            for cte_cols in cte_dict.values() {
                if let Some(cte_sources) = cte_cols.get(&col_upper).or_else(|| {
                    cte_cols
                        .iter()
                        .find(|(k, _)| k.eq_ignore_ascii_case(target_col))
                        .map(|(_, v)| v)
                }) && !cte_sources.is_empty()
                {
                    cte_source_sets.push(cte_sources);
                }
            }

            if cte_source_sets.is_empty() {
                continue;
            }

            // All CTE entries must agree on the source set
            let first = cte_source_sets[0];
            let all_agree = cte_source_sets.iter().all(|s| *s == first);
            if !all_agree || !first.is_subset(sources) || first.len() >= sources.len() {
                continue;
            }

            // Collect tables referenced by CTE dict sources
            let cte_tables: BTreeSet<String> = first
                .iter()
                .filter_map(|s| {
                    let parts: Vec<&str> = s.split('.').collect();
                    if parts.len() >= 2 {
                        Some(parts[0].trim_matches('"').to_uppercase())
                    } else {
                        None
                    }
                })
                .collect();

            // Only remove extras from SAME tables as CTE dict sources.
            // Extras from different tables may be from UNION/JOIN — keep them.
            let extras: Vec<String> = sources
                .difference(first)
                .filter(|s| {
                    let parts: Vec<&str> = s.split('.').collect();
                    if parts.len() >= 2 {
                        let table = parts[0].trim_matches('"').to_uppercase();
                        cte_tables.contains(&table)
                    } else {
                        false
                    }
                })
                .cloned()
                .collect();

            if !extras.is_empty() && extras.len() < sources.len() {
                for e in &extras {
                    sources.remove(e);
                }
            }
        }
    } // end if !inject_shared

    // ---------------------------------------------------------------
    // Filter output columns whose sources are ALL unresolvable.
    // ---------------------------------------------------------------
    // When every source reference for an output column points to a table
    // NOT in the provided schema (e.g., LATERAL FLATTEN aliases, CTE names
    // that couldn't be resolved, Hightouch metadata tables), the column is
    // likely synthetic/generated and Snowflake's ACCESS_USAGE won't track it.
    // Remove these output columns entirely.
    //
    // Only apply when schema is non-empty (otherwise we'd remove everything).
    if !flat_schema.is_empty() {
        let unresolvable_cols: Vec<String> = column_dict
            .iter()
            .filter(|(_, sources)| {
                !sources.is_empty()
                    && sources.iter().all(|ref_str| {
                        let parts: Vec<&str> = ref_str.split('.').collect();
                        if parts.len() < 2 {
                            return true; // no table part → unresolvable
                        }
                        // Extract table part (everything before the last dot-segment)
                        let table_part = parts[..parts.len() - 1]
                            .iter()
                            .map(|p| p.trim_matches('"'))
                            .collect::<Vec<_>>()
                            .join(".");
                        find_schema_table(&flat_schema, &table_part).is_none()
                    })
            })
            .map(|(k, _)| k.clone())
            .collect();
        for k in &unresolvable_cols {
            column_dict.remove(k);
        }
    }

    // ---------------------------------------------------------------
    // Normalize over-qualified source refs (strip db/schema prefix)
    // ---------------------------------------------------------------
    // Polyglot-sql can produce multi-part table names like "DW.MEMBER_PROFILE"
    // from SQL like `FROM "DW"."MEMBER_PROFILE"`. When the flat schema only
    // has "MEMBER_PROFILE" (no dots), strip the prefix to produce 2-part refs
    // like `"MEMBER_PROFILE"."COL"` instead of `"DW"."MEMBER_PROFILE"."COL"`.
    normalize_source_ref_prefixes(&mut column_dict, &flat_schema, dialect);
    if let Some(ref mut snap) = direct_snapshot {
        normalize_source_ref_prefixes(snap, &flat_schema, dialect);
    }

    // ---------------------------------------------------------------
    // Resolve LATERAL FLATTEN aliases in final output
    // ---------------------------------------------------------------
    // Replace any remaining refs to FLATTEN virtual aliases (e.g., "F"."VALUE")
    // with the actual column(s) being flattened.
    if !flatten_inputs.is_empty() {
        let alias_map_for_resolve = info
            .as_ref()
            .map(|i| &i.alias_map)
            .cloned()
            .unwrap_or_default();
        resolve_flatten_refs_in_dict(
            &mut column_dict,
            &flatten_inputs,
            &alias_map_for_resolve,
            &flat_schema,
            &cte_dict,
            dialect,
        );
        if let Some(ref mut snap) = direct_snapshot {
            resolve_flatten_refs_in_dict(
                snap,
                &flatten_inputs,
                &alias_map_for_resolve,
                &flat_schema,
                &cte_dict,
                dialect,
            );
        }
    }

    // ---------------------------------------------------------------
    // Default database/schema qualification
    // ---------------------------------------------------------------
    let needs_qualify = default_database.is_some() || default_schema.is_some();
    if needs_qualify {
        let qualify_ref = |source: &str| -> String {
            let dot_count = source.matches('.').count();
            // "table"."col" has 1 dot → unqualified table. 2+ dots → already qualified.
            if dot_count >= 2 {
                return source.to_string();
            }
            let mut prefix = Vec::new();
            if let Some(db) = default_database {
                let db_fmt = match dialect {
                    SqlDialect::Snowflake => format!("\"{}\"", db.to_uppercase()),
                    _ => format!("\"{}\"", db),
                };
                prefix.push(db_fmt);
            }
            if let Some(s) = default_schema {
                let s_fmt = match dialect {
                    SqlDialect::Snowflake => format!("\"{}\"", s.to_uppercase()),
                    _ => format!("\"{}\"", s),
                };
                prefix.push(s_fmt);
            }
            format!("{}.{}", prefix.join("."), source)
        };

        // Qualify column_dict source refs
        for sources in column_dict.values_mut() {
            let qualified: BTreeSet<String> = sources.iter().map(|s| qualify_ref(s)).collect();
            *sources = qualified;
        }
        // Also qualify the snapshot for correct direct/indirect classification
        if let Some(ref mut snap) = direct_snapshot {
            for sources in snap.values_mut() {
                let qualified: BTreeSet<String> = sources.iter().map(|s| qualify_ref(s)).collect();
                *sources = qualified;
            }
        }
        // Qualify shared_deps
        shared_deps = shared_deps.iter().map(|s| qualify_ref(s)).collect();
        // Qualify source_tables (unquoted format: "table" → "db.schema.table")
        source_tables = source_tables
            .iter()
            .map(|t| {
                if t.contains('.') {
                    t.clone()
                } else {
                    let mut parts = Vec::new();
                    if let Some(db) = default_database {
                        parts.push(match dialect {
                            SqlDialect::Snowflake => db.to_uppercase(),
                            _ => db.to_string(),
                        });
                    }
                    if let Some(s) = default_schema {
                        parts.push(match dialect {
                            SqlDialect::Snowflake => s.to_uppercase(),
                            _ => s.to_string(),
                        });
                    }
                    parts.push(t.clone());
                    parts.join(".")
                }
            })
            .collect();
    }

    // ---------------------------------------------------------------
    // Build flat lineage list with direct/indirect + lens metadata
    // ---------------------------------------------------------------

    // Pre-compute lens info from polyglot edges: (target, source_ref) → lens.
    // One pass over edges, O(E). Only includes non-CTE terminal edges.
    // Source refs are qualified to match column_dict keys when defaults are set.
    let mut edge_lens: HashMap<(String, String), (LensClassification, Vec<LensStep>)> =
        HashMap::with_capacity(base.edges.len());
    for edge in &base.edges {
        if edge.terminal_kind == TerminalKind::Table {
            let t = normalize_column_name(&edge.target_column, dialect);
            let mut s = edge_to_qualified_ref(edge, dialect);
            if needs_qualify && s.matches('.').count() < 2 {
                let mut prefix = Vec::new();
                if let Some(db) = default_database {
                    prefix.push(match dialect {
                        SqlDialect::Snowflake => format!("\"{}\"", db.to_uppercase()),
                        _ => format!("\"{}\"", db),
                    });
                }
                if let Some(sc) = default_schema {
                    prefix.push(match dialect {
                        SqlDialect::Snowflake => format!("\"{}\"", sc.to_uppercase()),
                        _ => format!("\"{}\"", sc),
                    });
                }
                if !prefix.is_empty() {
                    s = format!("{}.{}", prefix.join("."), s);
                }
            }
            let cls = polyglot_lens_to_classification(&edge.lens_type);
            let steps = polyglot_lens_code_to_steps(&edge.lens_code);
            edge_lens.insert((t, s), (cls, steps));
        }
    }

    // Single pass over column_dict to build both the flat list and the
    // HashMap output. Avoids iterating twice.
    let mut column_lineage: Vec<LineageEntry> = Vec::with_capacity(column_dict.len() * 4); // rough estimate
    let mut column_dict_out: HashMap<String, Vec<String>> =
        HashMap::with_capacity(column_dict.len());

    for (target, sources) in &column_dict {
        let sources_vec: Vec<String> = sources.iter().cloned().collect();
        for source in &sources_vec {
            // Classify: was this source present before shared injection?
            let is_direct = match &direct_snapshot {
                Some(snap) => snap.get(target).is_some_and(|pre| pre.contains(source)),
                None => true, // projection mode: everything is direct
            };
            let is_indirect = inject_shared && shared_deps.contains(source);

            if is_direct {
                let key = (target.clone(), source.clone());
                let (cls, steps) = edge_lens
                    .remove(&key)
                    .unwrap_or((LensClassification::Unknown, Vec::new()));
                column_lineage.push(LineageEntry {
                    source: source.clone(),
                    target: target.clone(),
                    lineage_type: LineageType::Direct,
                    lens_type: cls,
                    lens_code: steps,
                });
            }
            if is_indirect {
                column_lineage.push(LineageEntry {
                    source: source.clone(),
                    target: target.clone(),
                    lineage_type: LineageType::Indirect,
                    lens_type: LensClassification::Indirect,
                    lens_code: Vec::new(),
                });
            }
            // Edge case: source in final dict but not direct and not indirect
            // (shouldn't happen, but handle gracefully)
            if !is_direct && !is_indirect {
                column_lineage.push(LineageEntry {
                    source: source.clone(),
                    target: target.clone(),
                    lineage_type: LineageType::Direct,
                    lens_type: LensClassification::Unknown,
                    lens_code: Vec::new(),
                });
            }
        }
        column_dict_out.insert(target.clone(), sources_vec);
    }

    let errors: Vec<String> = base.errors.iter().map(|e| e.message.clone()).collect();

    Ok(CompleteLineageResult {
        column_dict: column_dict_out,
        column_lineage,
        source_tables: source_tables.into_iter().collect(),
        errors,
        default_database: default_database.map(String::from),
        default_schema: default_schema.map(String::from),
    })
}

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

    fn empty_schema() -> HashMap<String, serde_json::Value> {
        HashMap::new()
    }

    #[test]
    fn basic_select() {
        let opts = ColumnLineageOptions {
            inject_shared: true,
            ..Default::default()
        };
        let result = column_lineage(
            "SELECT a, b FROM table1",
            SqlDialect::Generic,
            &empty_schema(),
            Some(&opts),
        )
        .unwrap();

        assert!(result.column_dict.contains_key("a"));
        assert!(result.column_dict.contains_key("b"));
    }

    #[test]
    fn empty_sql() {
        let opts = ColumnLineageOptions {
            inject_shared: true,
            ..Default::default()
        };
        let result = column_lineage("", SqlDialect::Generic, &empty_schema(), Some(&opts)).unwrap();
        assert!(result.column_dict.is_empty());
    }

    #[test]
    fn comment_stripping() {
        let sql = "-- comment\nSELECT a /* inline */ FROM table1";
        let opts = ColumnLineageOptions {
            inject_shared: true,
            ..Default::default()
        };
        let result =
            column_lineage(sql, SqlDialect::Generic, &empty_schema(), Some(&opts)).unwrap();
        assert!(result.column_dict.contains_key("a"));
    }

    #[test]
    fn where_injection() {
        let mut schema = HashMap::new();
        schema.insert(
            "table1".to_string(),
            serde_json::json!({"a": "int", "b": "int", "c": "int"}),
        );
        let opts = ColumnLineageOptions {
            inject_shared: true,
            ..Default::default()
        };
        let result = column_lineage(
            "SELECT a, b FROM table1 WHERE c > 10",
            SqlDialect::Generic,
            &schema,
            Some(&opts),
        )
        .unwrap();

        let a_deps = result.column_dict.get("a").unwrap();
        let b_deps = result.column_dict.get("b").unwrap();
        assert!(a_deps.iter().any(|d| d.contains("\"c\"")));
        assert!(b_deps.iter().any(|d| d.contains("\"c\"")));
    }

    #[test]
    fn column_alias_through_subquery() {
        // Tests that column aliases are traced back to base table columns
        let mut schema = HashMap::new();
        schema.insert(
            "t1".to_string(),
            serde_json::json!({"a": "int", "b": "int"}),
        );
        let sql = "SELECT x FROM (SELECT a AS x FROM t1) sq";
        let result = column_lineage(sql, SqlDialect::Generic, &schema, None).unwrap();
        let x_deps = result.column_dict.get("x").unwrap();
        assert!(
            x_deps.iter().any(|d| d.contains("\"a\"")),
            "Expected x to trace back to t1.a, got: {:?}",
            x_deps
        );
    }

    #[test]
    fn column_alias_nested_subquery() {
        // Tests nested subquery alias tracing (_ht_to_filter pattern)
        let mut schema = HashMap::new();
        schema.insert(
            "base_table".to_string(),
            serde_json::json!({"col_a": "int", "col_b": "int"}),
        );
        let sql = r#"SELECT * FROM (
            SELECT sq.alias_col, sq.col_b
            FROM (
                SELECT col_a AS alias_col, col_b FROM base_table
            ) sq
        ) outer_sq"#;
        let result = column_lineage(sql, SqlDialect::Generic, &schema, None).unwrap();
        if let Some(deps) = result.column_dict.get("alias_col") {
            assert!(
                deps.iter().any(|d| d.contains("\"col_a\"")),
                "Expected alias_col to trace to base_table.col_a, got: {:?}",
                deps
            );
        }
    }

    #[test]
    fn debug_ht_to_filter_pattern() {
        // Reproduce the _ht_to_filter production pattern
        let mut schema = HashMap::new();
        schema.insert(
            "SHOPPING_ORDER_CUBE_V6".to_string(),
            serde_json::json!({"GUID": "VARCHAR", "STORE_ID": "VARCHAR"}),
        );
        let sql = r#"SELECT * FROM (
          SELECT
      "_ht_to_filter"."SINGLE_PURCHASE_MERCHANT"
,      "_ht_to_filter"."GUID"
    FROM
    (
      select guid,
  array_agg(store_id) as single_purchase_merchant
  from SHOPPING_ORDER_CUBE_V6
where store_id in (21991)
group by 1
  ) AS "_ht_to_filter"
        ) AS _ht_to_filter_pk
        WHERE
          "GUID"
          IS NOT NULL"#;

        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();

        // SINGLE_PURCHASE_MERCHANT should trace to STORE_ID
        let deps = result
            .column_dict
            .get("SINGLE_PURCHASE_MERCHANT")
            .expect("SINGLE_PURCHASE_MERCHANT should be in column_dict");
        let has_store_id = deps.iter().any(|d| d.contains("STORE_ID"));
        assert!(
            has_store_id,
            "SINGLE_PURCHASE_MERCHANT should trace to STORE_ID, got: {:?}\nFull column_dict: {:?}",
            deps, result.column_dict
        );
    }

    #[test]
    fn cte_alias_leak_prevention() {
        // Verifies that CTE/subquery names with empty source resolutions
        // don't leak as raw source references in lineage output.
        let mut schema = HashMap::new();
        schema.insert(
            "BASE_TABLE".to_string(),
            serde_json::json!({"COL_A": "int", "COL_B": "int"}),
        );

        // CTE chain: inner_cte -> filtered_cte, where filtered_cte
        // selects COL_C (not in schema). The CASE references COL_C
        // which can't be traced to a base table column.
        let sql = r#"SELECT * FROM (
          SELECT "_ht_to_filter"."RESULT_COL"
          FROM (
            WITH inner_cte AS (
              SELECT col_a, col_b FROM base_table
            ),
            filtered_cte AS (
              SELECT col_a, col_b FROM inner_cte
            )
            SELECT filtered_cte.col_a AS RESULT_COL
            FROM filtered_cte
          ) AS "_ht_to_filter"
        ) AS _ht_to_filter_pk"#;

        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();

        // RESULT_COL should trace to BASE_TABLE.COL_A, not to any CTE name
        if let Some(deps) = result.column_dict.get("RESULT_COL") {
            for dep in deps {
                assert!(
                    !dep.contains("inner_cte")
                        && !dep.contains("filtered_cte")
                        && !dep.contains("INNER_CTE")
                        && !dep.contains("FILTERED_CTE"),
                    "CTE name should not leak as source: {dep}"
                );
            }
        }
    }

    #[test]
    fn cte_chain_resolution() {
        // CTE chain: plan_old_simulated -> plan_old_simulated_raw -> plan_old_plan -> base table
        let mut schema = HashMap::new();
        schema.insert(
            "BASE_TABLE".to_string(),
            serde_json::json!({"GUID": "VARCHAR", "EMAIL": "VARCHAR", "_HT_OP_TYPE": "VARCHAR"}),
        );
        let sql = r#"
            WITH plan_old_plan AS (
                SELECT * FROM BASE_TABLE
            ), plan_old_simulated_raw AS (
                SELECT *, false AS _ht_prev_failed FROM plan_old_plan
            ), plan_old_simulated AS (
                SELECT *, ROW_NUMBER() OVER (ORDER BY GUID) AS _ht_idx FROM plan_old_simulated_raw
            )
            SELECT
                plan_old_simulated.GUID,
                plan_old_simulated.EMAIL,
                plan_old_simulated._HT_OP_TYPE
            FROM plan_old_simulated
        "#;
        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();
        // GUID should trace back to BASE_TABLE.GUID through the CTE chain
        let guid_deps = result
            .column_dict
            .get("GUID")
            .expect("GUID should be in output");
        assert!(
            guid_deps.iter().any(|d| d.contains("BASE_TABLE")),
            "GUID should trace to BASE_TABLE, got: {:?}\nFull: {:?}",
            guid_deps,
            result.column_dict
        );
    }

    #[test]
    fn cte_chain_with_union_and_join() {
        // Closer to Hightouch planner pattern
        let mut schema = HashMap::new();
        schema.insert(
            "PLAN_TABLE".to_string(),
            serde_json::json!({"GUID": "VARCHAR", "_HT_OP_TYPE": "VARCHAR", "_HT_PLAN_INDEX": "INT"}),
        );
        schema.insert(
            "REJECTIONS_TABLE".to_string(),
            serde_json::json!({"ID": "VARCHAR", "_HT_OP_TYPE": "VARCHAR"}),
        );
        schema.insert(
            "NEW_RESULTS".to_string(),
            serde_json::json!({"GUID": "VARCHAR", "EMAIL": "VARCHAR"}),
        );

        let sql = r#"
            WITH plan_old_plan AS (
                SELECT * FROM PLAN_TABLE
            ), plan_old_rejections AS (
                SELECT DISTINCT id, _ht_op_type FROM REJECTIONS_TABLE
            ), plan_old_simulated_raw AS (
                SELECT *, false AS _ht_prev_failed_simulated FROM plan_old_plan AS _ht_old_plan
                WHERE _ht_op_type = 'unchanged'
                UNION ALL
                SELECT _ht_old_plan.*,
                    _ht_old_rejections.id IS NOT NULL AS _ht_prev_failed_simulated
                FROM plan_old_plan AS _ht_old_plan
                LEFT JOIN plan_old_rejections AS _ht_old_rejections
                    ON _ht_old_plan.GUID = _ht_old_rejections.ID
                WHERE _ht_old_plan._ht_op_type = 'added'
            ), plan_old_simulated AS (
                SELECT *, ROW_NUMBER() OVER (ORDER BY GUID) AS _ht_distinct_pk_index
                FROM plan_old_simulated_raw AS _ht_old_simulated_raw
            ), plan_new_results AS (
                SELECT * FROM NEW_RESULTS
            ), plan_outer_join AS (
                SELECT
                    plan_old_simulated._ht_prev_failed_simulated AS _ht_prev_failed,
                    plan_old_simulated._ht_op_type AS _ht_old_op_type,
                    plan_old_simulated.GUID AS _ht_oldguid,
                    plan_new_results.GUID,
                    plan_new_results.EMAIL
                FROM plan_old_simulated
                FULL OUTER JOIN plan_new_results
                    ON plan_old_simulated.GUID = plan_new_results.GUID
            )
            SELECT * FROM plan_outer_join
        "#;

        // Check raw edges
        let base = super::polyglot_column_lineage(sql, SqlDialect::Snowflake).unwrap();
        assert!(
            !base.edges.is_empty(),
            "Expected non-empty edges from polyglot"
        );

        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();

        // _HT_OLDGUID should trace to PLAN_TABLE.GUID
        let oldguid = result.column_dict.get("_HT_OLDGUID");
        assert!(
            oldguid.map_or(false, |d| d.iter().any(|s| s.contains("PLAN_TABLE"))),
            "_HT_OLDGUID should trace to PLAN_TABLE, got: {:?}",
            oldguid
        );

        // EMAIL should trace to NEW_RESULTS.EMAIL
        let email = result.column_dict.get("EMAIL");
        assert!(
            email.map_or(false, |d| d.iter().any(|s| s.contains("NEW_RESULTS"))),
            "EMAIL should trace to NEW_RESULTS, got: {:?}",
            email
        );
    }

    #[test]
    fn flatten_schema_flat() {
        let schema: HashMap<String, serde_json::Value> =
            serde_json::from_str(r#"{"table1": {"a": "int", "b": "int"}}"#).unwrap();
        let flat = flatten_schema(&schema);
        assert!(flat.contains_key("table1"));
        assert_eq!(flat["table1"].len(), 2);
    }

    #[test]
    fn flatten_schema_nested() {
        let schema: HashMap<String, serde_json::Value> =
            serde_json::from_str(r#"{"db": {"schema": {"table1": {"a": "int", "b": "int"}}}}"#)
                .unwrap();
        let flat = flatten_schema(&schema);
        assert!(flat.contains_key("db.schema.table1"));
        assert_eq!(flat["db.schema.table1"].len(), 2);
    }

    #[test]
    fn select_star_with_expression_alias() {
        // Tests: SELECT * from subquery that has expression aliases
        let mut schema = HashMap::new();
        schema.insert(
            "MEMBER_BONUS".to_string(),
            serde_json::json!({
                "MEMBER_BONUS_ID": "VARCHAR",
                "QUALIFYING_ORDERS": "VARCHAR",
                "MODIFIED_BY": "VARCHAR",
                "LAST_MODIFIED_DATE": "VARCHAR"
            }),
        );
        // Inner SELECT casts + renames: `value::int AS qualifying_order_id`
        // This is an expression alias. Outer SELECT * should still produce it.
        let sql = r#"select *
from (
  select member_bonus_id, modified_by, last_modified_date
  from MEMBER_BONUS
  where tenant = 'ebates.com'
) A
group by all"#;

        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();
        assert!(
            result.column_dict.contains_key("MEMBER_BONUS_ID"),
            "Missing MEMBER_BONUS_ID. Keys: {:?}",
            result.column_dict.keys().collect::<Vec<_>>()
        );
        assert!(result.column_dict.contains_key("MODIFIED_BY"));
        assert!(result.column_dict.contains_key("LAST_MODIFIED_DATE"));
    }

    #[test]
    fn select_star_from_union_subquery() {
        // SELECT * FROM (... UNION ALL ...) should expand to all output columns
        let mut schema = HashMap::new();
        schema.insert(
            "MEMBER_BONUS".to_string(),
            serde_json::json!({
                "MODIFIED_BY": "VARCHAR",
                "MEMBER_BONUS_ID": "VARCHAR",
                "QUALIFYING_ORDER_ID": "VARCHAR",
                "LAST_MODIFIED_DATE": "VARCHAR"
            }),
        );
        let sql = r#"SELECT * FROM (
            SELECT member_bonus_id, qualifying_order_id, modified_by, last_modified_date
            FROM MEMBER_BONUS WHERE tenant = 'ebates.com'
            UNION ALL
            SELECT member_bonus_id, qualifying_order_id, modified_by, last_modified_date
            FROM MEMBER_BONUS WHERE state_id = 50
        ) A GROUP BY ALL"#;

        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();
        // All 4 columns should be present
        assert!(
            result.column_dict.contains_key("MEMBER_BONUS_ID"),
            "Missing MEMBER_BONUS_ID. Keys: {:?}",
            result.column_dict.keys().collect::<Vec<_>>()
        );
        assert!(result.column_dict.contains_key("QUALIFYING_ORDER_ID"));
        assert!(result.column_dict.contains_key("MODIFIED_BY"));
        assert!(result.column_dict.contains_key("LAST_MODIFIED_DATE"));
    }

    #[test]
    fn snowflake_variant_access_lineage() {
        // Snowflake variant syntax: col:field::TYPE should trace to the base column
        let mut schema = HashMap::new();
        schema.insert(
            "GLOBAL_EMAIL_HISTORY".to_string(),
            serde_json::json!({"MEMBER_ID": "VARCHAR", "ATTRIBUTES": "VARCHAR"}),
        );
        let sql = r#"SELECT member_id AS user_id, ATTRIBUTES:MERCHANT_ID::String AS merchant_id FROM SUMMARY.GLOBAL_EMAIL_HISTORY"#;
        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();

        let merchant = result
            .column_dict
            .get("MERCHANT_ID")
            .expect("MERCHANT_ID missing");
        assert!(
            merchant.iter().any(|s| s.contains("ATTRIBUTES")),
            "MERCHANT_ID should trace to ATTRIBUTES column, got: {merchant:?}"
        );
        let user = result.column_dict.get("USER_ID").expect("USER_ID missing");
        assert!(
            user.iter().any(|s| s.contains("MEMBER_ID")),
            "USER_ID should trace to MEMBER_ID, got: {user:?}"
        );
    }

    #[test]
    fn snowflake_object_construct_star() {
        // OBJECT_CONSTRUCT(*) should expand * to all table columns
        let mut schema = HashMap::new();
        schema.insert(
            "FEED_LOG".to_string(),
            serde_json::json!({
                "GUID": "VARCHAR", "AMOUNT": "VARCHAR", "TENANT": "VARCHAR"
            }),
        );
        let sql = r#"SELECT guid AS KEY, object_construct(*) AS value FROM ASSET.FEED_LOG"#;
        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();

        let value = result.column_dict.get("VALUE").expect("VALUE missing");
        assert!(
            value.iter().any(|s| s.contains("GUID")),
            "VALUE should include GUID, got: {value:?}"
        );
        assert!(
            value.iter().any(|s| s.contains("AMOUNT")),
            "VALUE should include AMOUNT, got: {value:?}"
        );
        assert!(
            value.iter().any(|s| s.contains("TENANT")),
            "VALUE should include TENANT, got: {value:?}"
        );
    }

    #[test]
    fn ambiguous_column_resolves_to_from_table() {
        // Two tables with identical columns; only one in FROM clause.
        // Engine should resolve to the FROM-clause table, not the other.
        let mut schema = HashMap::new();
        schema.insert(
            "APPDB_CURRENCY_EXCHANGE_RATE".to_string(),
            serde_json::json!({"EURO_EXCHANGE_RATE": "VARCHAR", "EXCHANGE_CURRENCY": "VARCHAR", "EXCHANGE_DATE": "VARCHAR", "USD_EXCHANGE_RATE": "VARCHAR"}),
        );
        schema.insert(
            "RR_APPDB_CURRENCY_EXCHANGE_RATE".to_string(),
            serde_json::json!({"EURO_EXCHANGE_RATE": "VARCHAR", "EXCHANGE_CURRENCY": "VARCHAR", "EXCHANGE_DATE": "VARCHAR", "USD_EXCHANGE_RATE": "VARCHAR"}),
        );

        let sql = r#"SELECT EXCHANGE_DATE AS conversion_date,
            EURO_EXCHANGE_RATE AS rate,
            USD_EXCHANGE_RATE AS usd_rate,
            EXCHANGE_CURRENCY AS currency
        FROM DW_SS.RR_APPDB_CURRENCY_EXCHANGE_RATE
        WHERE INSERT_TIMESTAMP >= '2026-02-18 12:14:17.518'
        QUALIFY ROW_NUMBER() OVER (PARTITION BY EXCHANGE_CURRENCY, EXCHANGE_DATE ORDER BY insert_timestamp desc) = 1"#;

        let result = column_lineage(sql, SqlDialect::Snowflake, &schema, None).unwrap();

        // All sources should come from RR_APPDB_CURRENCY_EXCHANGE_RATE (the FROM table)
        let correct = "RR_APPDB_CURRENCY_EXCHANGE_RATE";
        let wrong = "\"APPDB_CURRENCY_EXCHANGE_RATE\"";
        for (col, srcs) in &result.column_dict {
            for src in srcs {
                // Check that the source starts with the correct table (not the wrong one)
                let is_wrong = src.starts_with(wrong) && !src.contains("RR_APPDB");
                assert!(
                    !is_wrong,
                    "Column {} resolved to wrong table: {} (should be {})",
                    col, src, correct
                );
            }
        }
    }

    #[test]
    fn flat_lineage_direct_and_indirect() {
        let mut schema = HashMap::new();
        schema.insert(
            "t1".to_string(),
            serde_json::json!({"a": "int", "b": "int", "c": "int"}),
        );
        let opts = ColumnLineageOptions {
            inject_shared: true,
            ..Default::default()
        };
        let result = column_lineage(
            "SELECT a, b FROM t1 WHERE c > 10",
            SqlDialect::Generic,
            &schema,
            Some(&opts),
        )
        .unwrap();

        // Flat list should have entries
        assert!(
            !result.column_lineage.is_empty(),
            "column_lineage should not be empty"
        );

        // Direct links: a←t1.a, b←t1.b
        let direct: Vec<_> = result
            .column_lineage
            .iter()
            .filter(|e| e.lineage_type == LineageType::Direct)
            .collect();
        assert!(
            direct
                .iter()
                .any(|e| e.target == "a" && e.source.contains("\"a\"")),
            "Missing direct link a←t1.a: {:?}",
            direct
        );
        assert!(
            direct
                .iter()
                .any(|e| e.target == "b" && e.source.contains("\"b\"")),
            "Missing direct link b←t1.b: {:?}",
            direct
        );

        // Indirect links: c used in WHERE → both a and b get it
        let indirect: Vec<_> = result
            .column_lineage
            .iter()
            .filter(|e| e.lineage_type == LineageType::Indirect)
            .collect();
        let indirect_targets: std::collections::BTreeSet<_> =
            indirect.iter().map(|e| e.target.as_str()).collect();
        assert!(
            indirect_targets.contains("a") && indirect_targets.contains("b"),
            "WHERE column c should appear as indirect for both a and b, got targets: {:?}",
            indirect_targets
        );
        // All indirect entries should have Indirect lens_type
        for e in &indirect {
            assert_eq!(e.lens_type, LensClassification::Indirect);
        }
    }

    #[test]
    fn flat_lineage_projection_mode() {
        let mut schema = HashMap::new();
        schema.insert(
            "t1".to_string(),
            serde_json::json!({"a": "int", "c": "int"}),
        );
        let result = column_lineage(
            "SELECT a FROM t1 WHERE c > 10",
            SqlDialect::Generic,
            &schema,
            None,
        )
        .unwrap();

        // In projection mode, all entries should be Direct
        for e in &result.column_lineage {
            assert_eq!(
                e.lineage_type,
                LineageType::Direct,
                "Projection mode should only have Direct links, got: {:?}",
                e
            );
        }
        // WHERE column c should NOT appear in projection mode
        assert!(
            !result
                .column_lineage
                .iter()
                .any(|e| e.source.contains("\"c\"")),
            "Projection mode should not include WHERE column c"
        );
    }

    #[test]
    fn default_database_schema_qualification() {
        let mut schema = HashMap::new();
        schema.insert(
            "t1".to_string(),
            serde_json::json!({"a": "int", "b": "int"}),
        );
        let opts = ColumnLineageOptions {
            default_database: Some("mydb".to_string()),
            default_schema: Some("public".to_string()),
            inject_shared: false,
        };
        let result = column_lineage(
            "SELECT a FROM t1",
            SqlDialect::Snowflake,
            &schema,
            Some(&opts),
        )
        .unwrap();

        // Source refs should be fully qualified
        let a_deps = result.column_dict.get("A").unwrap();
        assert!(
            a_deps.iter().any(|d| d.contains("\"MYDB\".\"PUBLIC\"")),
            "Expected qualified refs with MYDB.PUBLIC, got: {:?}",
            a_deps
        );
        // Result should report the defaults
        assert_eq!(result.default_database.as_deref(), Some("mydb"));
        assert_eq!(result.default_schema.as_deref(), Some("public"));
        // Source tables qualified when populated
        for t in &result.source_tables {
            if !t.contains('.') {
                panic!("Unqualified source table: {t}");
            }
        }
        // Flat lineage links should also be qualified
        for entry in &result.column_lineage {
            assert!(
                entry.source.contains("\"MYDB\""),
                "Lineage link source should be qualified: {:?}",
                entry
            );
        }
    }
}
