"""Get column-level lineage through SQL transformations.

This module implements the get_column_lineage tool for dbt Core MCP.
Uses sqlglot to parse SQL and trace column dependencies through CTEs and transformations.

Architecture:
- Unified SQL parsing approach for both upstream and downstream directions
- Two-stage process: (1) resolve output columns, (2) analyze lineage per column
- Direction-agnostic helpers: _prepare_model_analysis, _analyze_column_lineage
- Consistent fallback order: SQL-derived → warehouse → manifest → none
"""

import logging
from typing import Any

from fastmcp.dependencies import Depends  # type: ignore[reportAttributeAccessIssue]
from fastmcp.server.context import Context
from sqlglot import exp, parse_one
from sqlglot.errors import SqlglotError
from sqlglot.lineage import lineage
from sqlglot.optimizer.scope import build_scope

from ..context import DbtCoreServerContext
from ..dbt.manifest import ManifestLoader
from ..dependencies import get_state
from . import dbtTool

logger = logging.getLogger(__name__)


def _wrap_final_select(sql: str, column_name: str, dialect: str = "databricks") -> exp.Expression:
    """Wrap the final SELECT in a CTE to enable lineage tracing through SELECT *.

    Transforms:
        SELECT * FROM final
    Into:
        WITH __lineage_final__ AS (SELECT * FROM final)
        SELECT column_name FROM __lineage_final__

    This allows lineage() to trace specific columns through schemaless queries.

    Args:
        sql: SQL query to wrap
        column_name: Column to trace
        dialect: SQL dialect

    Returns:
        Modified AST with wrapped final SELECT
    """
    ast = parse_one(sql, dialect=dialect)

    # Get the root SELECT
    root_select = ast if isinstance(ast, exp.Select) else ast.find(exp.Select)

    if not root_select:
        return ast

    # Save the WITH clause before clearing
    with_clause = root_select.args.get("with_")

    # Copy entire root SELECT and strip its WITH
    wrapper_select = root_select.copy()
    wrapper_select.set("with_", None)

    wrapper_cte = exp.CTE(this=wrapper_select, alias=exp.TableAlias(this=exp.Identifier(this="__lineage_final__")))

    # Clear ALL args from root
    for key in list(root_select.args.keys()):
        root_select.set(key, None)

    # Add wrapper to with_clause or create new one
    if with_clause:
        with_clause.expressions.append(wrapper_cte)
    else:
        with_clause = exp.With(expressions=[wrapper_cte])

    # Rebuild root SELECT with only what we need
    root_select.set("with_", with_clause)
    root_select.set("expressions", [exp.Column(this=exp.Identifier(this=column_name))])
    root_select.set("from_", exp.From(this=exp.Table(this=exp.Identifier(this="__lineage_final__"))))

    return ast


# Export for testing
__all__ = [
    "implementation",
    "get_column_lineage",
    "_prepare_model_analysis",
    "_analyze_column_lineage",
    "_build_schema_mapping",
    "_resolve_output_columns",
    "_resolve_wildcard_column_in_table",
    "_resolve_unresolved_table_reference",
    "_extract_transformations_with_sources",
    "_extract_dependencies_from_lineage",
    "_format_lineage_response",
    "_map_dbt_adapter_to_sqlglot_dialect",
]


def _map_dbt_adapter_to_sqlglot_dialect(adapter_type: str) -> str:
    """Map dbt adapter type to sqlglot dialect.

    Args:
        adapter_type: The dbt adapter type from manifest metadata

    Returns:
        sqlglot dialect name
    """
    # Direct matches between dbt adapter and sqlglot dialect
    ADAPTER_TO_DIALECT = {
        "athena": "athena",
        "bigquery": "bigquery",
        "clickhouse": "clickhouse",
        "databricks": "databricks",
        "doris": "doris",
        "dremio": "dremio",
        "duckdb": "duckdb",
        "fabric": "fabric",
        "hive": "hive",
        "materialize": "materialize",
        "mysql": "mysql",
        "oracle": "oracle",
        "postgres": "postgres",
        "postgresql": "postgres",  # Some adapters use postgresql
        "redshift": "redshift",
        "risingwave": "risingwave",
        "singlestore": "singlestore",
        "snowflake": "snowflake",
        "spark": "spark",
        "sqlite": "sqlite",
        "starrocks": "starrocks",
        "teradata": "teradata",
        "trino": "trino",
        # Adapters that need dialect mapping
        "synapse": "tsql",  # Azure Synapse uses T-SQL
        "sqlserver": "tsql",  # SQL Server uses T-SQL
        "glue": "spark",  # AWS Glue uses Spark
        "fabricspark": "spark",  # Fabric Lakehouse uses Spark
    }

    adapter_lower = adapter_type.lower()
    dialect = ADAPTER_TO_DIALECT.get(adapter_lower)

    if dialect:
        logger.debug(f"Mapped dbt adapter '{adapter_type}' to sqlglot dialect '{dialect}'")
        return dialect

    # Fallback: use adapter type as-is (might work for some cases)
    logger.warning(f"No explicit mapping for dbt adapter '{adapter_type}', using as-is for sqlglot")
    return adapter_lower


# ========== Unified Helpers (Direction-Agnostic) ==========
# These helpers are used by both upstream and downstream lineage analysis,
# ensuring consistent behavior regardless of direction:
#
# 1. _prepare_model_analysis(): Gets resource_info + compiled_sql + schema_mapping
# 2. _analyze_column_lineage(): Runs sqlglot lineage with consistent error handling
# 3. _resolve_output_columns(): Resolves output columns with fallback order
#
# This unified approach ensures:
# - Same SQL parsing logic for both directions
# - Consistent error handling and logging
# - No code duplication
# - Easier to maintain and test


def _build_schema_mapping(manifest: ManifestLoader, upstream_lineage: dict[str, Any]) -> dict[str, Any]:
    """Build schema mapping from upstream models for sqlglot.

    Args:
        manifest: ManifestLoader instance
        upstream_lineage: Upstream lineage dict from manifest.get_lineage()

    Returns:
        Schema mapping in format: {database: {schema: {table: {column: type}}}}
    """
    schema_mapping: dict[str, Any] = {}

    if "upstream" not in upstream_lineage:
        return schema_mapping

    for upstream_node in upstream_lineage["upstream"]:
        try:
            # Use node name, not unique_id (get_resource_info expects name)
            node_name = upstream_node.get("name")
            unique_id = upstream_node.get("unique_id")
            if not node_name or not unique_id:
                continue

            # Extract resource type from unique_id (e.g. "seed.project.name" -> "seed")
            resource_type = unique_id.split(".", 1)[0] if "." in unique_id else "model"

            node_info = manifest.get_resource_info(node_name, resource_type=resource_type, include_database_schema=True, include_compiled_sql=False)

            database = node_info.get("database", "").lower()
            schema = node_info.get("schema", "").lower()
            # For sources, use identifier; for models/seeds, use alias or name
            table = node_info.get("identifier") or node_info.get("alias") or node_info.get("name", "").lower()

            if database and schema and table:
                # Add columns with their types
                columns = node_info.get("database_columns", [])

                if not columns:
                    manifest_columns = node_info.get("columns", {})
                    if manifest_columns:
                        columns = {col_name: {"type": (col_info.get("data_type") or col_info.get("type") or "string")} for col_name, col_info in manifest_columns.items()}

                if not columns:
                    column_map = {col_info.get("col_name", "").lower(): col_info.get("type", "string").lower() for col_info in columns}
                else:
                    # Dict format: {"customer_id": {"type": "INTEGER"}}
                    column_map = {col_name.lower(): col_info.get("type", "string").lower() for col_name, col_info in columns.items()}

                if not column_map:
                    continue

                if database not in schema_mapping:
                    schema_mapping[database] = {}
                if schema not in schema_mapping[database]:
                    schema_mapping[database][schema] = {}

                schema_mapping[database][schema][table] = column_map
        except Exception as e:
            logger.warning(f"Could not load schema for upstream node {upstream_node.get('unique_id')}: {e}")
            continue

    return schema_mapping


def _find_table_columns(schema_mapping: dict[str, Any], table_name: str) -> list[str]:
    """Find column names for a table in the schema mapping.

    Args:
        schema_mapping: Schema mapping in format {db: {schema: {table: {column: type}}}}
        table_name: Table name to look up

    Returns:
        List of column names for the table
    """
    table_lower = table_name.lower()
    for database_mapping in schema_mapping.values():
        for schema_mapping_for_db in database_mapping.values():
            if table_lower in schema_mapping_for_db:
                return list(schema_mapping_for_db[table_lower].keys())
    return []


def _normalize_relation_name(value: str) -> str:
    """Normalize relation names for matching."""
    return value.replace('"', "").replace("`", "").replace("[", "").replace("]", "").strip().lower()


def _add_relation_keys(lookup: dict[str, str], database: str | None, schema: str | None, identifier: str | None, unique_id: str) -> None:
    if not identifier:
        return

    if database and schema:
        lookup[_normalize_relation_name(f"{database}.{schema}.{identifier}")] = unique_id

    if schema:
        lookup[_normalize_relation_name(f"{schema}.{identifier}")] = unique_id

    lookup[_normalize_relation_name(identifier)] = unique_id


def _build_relation_lookup(manifest: ManifestLoader) -> dict[str, str]:
    """Build relation name -> unique_id lookup from the manifest."""
    lookup: dict[str, str] = {}
    manifest_dict = manifest.get_manifest_dict()

    for node in manifest_dict.get("nodes", {}).values():
        if not isinstance(node, dict):
            continue

        unique_id = node.get("unique_id")
        if not unique_id:
            continue

        relation_name = node.get("relation_name")
        if relation_name:
            lookup[_normalize_relation_name(relation_name)] = unique_id

        database = node.get("database")
        schema = node.get("schema")
        identifier = node.get("alias") or node.get("name")
        _add_relation_keys(lookup, database, schema, identifier, unique_id)

    for source in manifest_dict.get("sources", {}).values():
        if not isinstance(source, dict):
            continue

        unique_id = source.get("unique_id")
        if not unique_id:
            continue

        relation_name = source.get("relation_name")
        if relation_name:
            lookup[_normalize_relation_name(relation_name)] = unique_id

        database = source.get("database")
        schema = source.get("schema")
        identifier = source.get("identifier") or source.get("name")
        _add_relation_keys(lookup, database, schema, identifier, unique_id)

    return lookup


def _get_output_columns_from_sql(compiled_sql: str, schema_mapping: dict[str, Any], dialect: str = "databricks") -> list[str]:
    """Extract output column names from compiled SQL.

    Handles SELECT * from a single CTE or table by expanding from schema mapping
    or the CTE's projections.

    Args:
        compiled_sql: Compiled SQL string
        schema_mapping: Schema mapping for table expansion
        dialect: SQL dialect for parsing (default: databricks)

    Returns:
        List of output column names
    """
    try:
        ast = parse_one(compiled_sql, dialect=dialect)
    except Exception:
        return []

    root_scope = build_scope(ast)
    if not root_scope:
        return []

    select = root_scope.expression if isinstance(root_scope.expression, exp.Select) else root_scope.expression.find(exp.Select)
    if not select:
        return []

    projections = list(select.expressions)
    if projections and all(isinstance(p, exp.Star) for p in projections):
        if len(root_scope.selected_sources) == 1:
            _, (_, source) = next(iter(root_scope.selected_sources.items()))
            if isinstance(source, exp.Table):
                return _find_table_columns(schema_mapping, source.name)

            # CTE or subquery scope
            if hasattr(source, "expression"):
                cte_select = source.expression if isinstance(source.expression, exp.Select) else source.expression.find(exp.Select)
                if cte_select:
                    return [proj.alias_or_name for proj in cte_select.expressions if proj.alias_or_name]
        return []

    return [proj.alias_or_name for proj in projections if proj.alias_or_name]


def _resolve_output_columns(
    compiled_sql: str,
    schema_mapping: dict[str, Any],
    resource_info: dict[str, Any],
    dialect: str = "databricks",
) -> tuple[dict[str, Any], str]:
    """Resolve output columns for a model using resource-specific strategies.

    For Sources & Seeds (Database Tables):
    1) Warehouse columns (database_columns) - database is truth
    2) Manifest columns (schema.yml) - fallback
    3) Wildcard "*" - table not built yet, but table reference known

    For Models (SQL Transformations):
    1) SQL-derived projections ONLY - compiled SQL is truth
       (No fallbacks - if SQL parsing fails, no columns available)

    Args:
        compiled_sql: Compiled SQL string
        schema_mapping: Schema mapping for table expansion
        resource_info: Resource information dict
        dialect: SQL dialect for parsing (default: databricks)

    Returns:
        (output_columns_dict, source_label)
    """
    output_columns_dict: dict[str, Any] = {}

    # For sources AND seeds: database-first with wildcard fallback
    if resource_info.get("resource_type") in ["source", "seed"]:
        database_columns = resource_info.get("database_columns", [])
        if isinstance(database_columns, list) and database_columns:
            output_columns_dict = {col.get("col_name", ""): {} for col in database_columns if col.get("col_name")}
            if output_columns_dict:
                return output_columns_dict, "warehouse"
        elif isinstance(database_columns, dict) and database_columns:
            return {col_name: {} for col_name in database_columns.keys()}, "warehouse"

        # Fallback to manifest columns
        manifest_columns = resource_info.get("columns", {})
        output_columns_dict = {col_name: {} for col_name in manifest_columns.keys()}
        if output_columns_dict:
            return output_columns_dict, "manifest"

        # Final fallback: table not built yet, return wildcard
        return {"*": {}}, "wildcard"

    # For models/seeds: SQL parsing ONLY (compiled SQL is truth)
    output_columns = _get_output_columns_from_sql(compiled_sql, schema_mapping, dialect)
    if output_columns:
        return {col: {} for col in output_columns}, "sql"

    return {}, "none"


def _resolve_wildcard_column_in_table(table_name: str, schema_mapping: dict[str, Any]) -> str | None:
    """Resolve wildcard (*) column to specific column name within a known table.

    When we have a resolved table but wildcard column (e.g., "column": "*", "table": "stg_customers"),
    this function attempts to find the specific column being traced in that table.

    Args:
        table_name: Name of the table (without quotes)
        schema_mapping: Schema mapping with table and column information

    Returns:
        Resolved column name if found, None otherwise
    """
    if table_name not in schema_mapping:
        return None

    table_info = schema_mapping[table_name]
    resource_type = table_info.get("resource_type", "")

    # For sources and seeds: use database-first approach
    if resource_type in ["source", "seed"]:
        # Check database columns first
        database_columns = table_info.get("database_columns", [])
        if isinstance(database_columns, list) and database_columns:
            # Return first database column as representative
            first_col = database_columns[0]
            if isinstance(first_col, dict):
                return first_col.get("col_name")
        elif isinstance(database_columns, dict) and database_columns:
            # Return first key from database columns
            return next(iter(database_columns.keys()))

        # Fallback to manifest columns
        manifest_columns = table_info.get("columns", {})
        if manifest_columns:
            return next(iter(manifest_columns.keys()))

    # For models: would use SQL parsing (handled elsewhere)
    return None


def _resolve_unresolved_table_reference(col_name: str, schema_mapping: dict[str, Any]) -> dict[str, Any] | None:
    """Resolve unresolved table references using schema mapping and database-first approach.

    When sqlglot can't resolve a table reference (returns None), this function
    attempts to find which table in the schema mapping contains the requested column,
    using database-first resolution for sources and seeds.

    Args:
        col_name: Column name to search for
        schema_mapping: Schema mapping with table and column information

    Returns:
        Dictionary with resolved table info if found, None otherwise:
        {
            "table_name": str,
            "database": str | None,
            "schema": str | None,
            "resolved_column": str | None  # If resolved from wildcard
        }
    """
    for table_name, table_info in schema_mapping.items():
        resource_type = table_info.get("resource_type", "")

        # For sources and seeds: use database-first approach
        if resource_type in ["source", "seed"]:
            # Check database columns first
            database_columns = table_info.get("database_columns", [])
            if isinstance(database_columns, list):
                for db_col in database_columns:
                    if isinstance(db_col, dict) and db_col.get("col_name", "").lower() == col_name.lower():
                        return {"table_name": f'"{table_name}"', "database": table_info.get("database"), "schema": table_info.get("schema"), "resolved_column": db_col.get("col_name")}
            elif isinstance(database_columns, dict):
                if col_name.lower() in [k.lower() for k in database_columns.keys()]:
                    # Find the actual key with proper case
                    actual_col = next(k for k in database_columns.keys() if k.lower() == col_name.lower())
                    return {"table_name": f'"{table_name}"', "database": table_info.get("database"), "schema": table_info.get("schema"), "resolved_column": actual_col}

            # Fallback to manifest columns for sources/seeds
            manifest_columns = table_info.get("columns", {})
            if col_name.lower() in [k.lower() for k in manifest_columns.keys()]:
                actual_col = next(k for k in manifest_columns.keys() if k.lower() == col_name.lower())
                return {"table_name": f'"{table_name}"', "database": table_info.get("database"), "schema": table_info.get("schema"), "resolved_column": actual_col}

    return None


def _extract_transformations_with_sources(lineage_node: Any) -> list[dict[str, Any]]:
    """Extract transformations with namespaced IDs, types, and source references.

    Uses a two-pass approach:
    1. First pass: Create all transformations and build lookup map
    2. Second pass: Extract sources by parsing expressions for CTE/table references

    Args:
        lineage_node: Result from sqlglot.lineage()

    Returns:
        List of transformation dicts with structure:
        {
            "id": "cte:name" or "table:name",  # Namespaced identifier
            "type": "cte" or "table",          # Type of transformation
            "column": "column_name",           # Column being transformed
            "expression": "...",               # Optional: how it's computed (if not pass-through)
            "sources": ["cte:other", ...]      # Upstream dependencies
        }
    """
    transform_map: dict[str, dict[str, Any]] = {}  # id -> transform dict
    cte_to_id: dict[str, str] = {}  # cte_name -> transform id (for lookup)
    nodes_with_data: list[tuple[str, Any, str]] = []  # (transform_id, node, expression)
    outer_query_sources: set[str] = set()  # Track what the outer query references
    union_branches: dict[str, list[dict[str, Any]]] = {}  # reference_node_name -> list of branch info

    # FIRST PASS: Create all transformations and build lookup map
    for node in lineage_node.walk():
        if not hasattr(node, "name") or not node.name:
            continue

        # Handle numeric nodes (UNION branches)
        if "." not in node.name:
            # Check if this is a UNION branch node (has reference_node_name attribute)
            if hasattr(node, "reference_node_name") and node.reference_node_name:
                ref_cte = node.reference_node_name

                # Extract expression and sources from this branch
                branch_info: dict[str, Any] = {}
                if hasattr(node, "expression") and node.expression:
                    expr_str = str(node.expression)
                    branch_info["expression"] = expr_str if len(expr_str) <= 200 else expr_str[:197] + "..."

                    # Extract column name (will be used for the union transformation)
                    # Expression is like "all_segments_union.customer_id AS customer_id"
                    if " AS " in expr_str:
                        branch_info["column"] = expr_str.split(" AS ")[-1].strip()

                    # Store full expression for sources extraction later
                    branch_info["_full_expr"] = expr_str

                # Group branches by their parent CTE
                union_branches.setdefault(ref_cte, []).append(branch_info)
            continue

        parts = node.name.split(".", 1)
        cte_or_table = parts[0]
        col_name = parts[1] if len(parts) > 1 else node.name

        # Special handling for wrapper CTE - extract what the outer query references
        if cte_or_table == "__lineage_final__":
            # Extract what __lineage_final__ references (could be CTE or table)
            if hasattr(node, "expression") and node.expression:
                expr_str = str(node.expression)
                # Expression will be like "final.customer_id" (CTE) or "si.gold_itemkey" (table alias)
                # Look for references in the expression
                for potential_ref in expr_str.split():
                    if "." in potential_ref:
                        # Could be "final.customer_id" or "si.gold_itemkey"
                        ref_name = potential_ref.split(".")[0].strip("(),")
                        if ref_name and not ref_name.isdigit():
                            # Try as CTE first, then as table
                            # We'll resolve the actual type in the second pass
                            outer_query_sources.add(ref_name)
            continue

        source_type = type(node.source).__name__ if hasattr(node, "source") else None

        # Get the actual CTE/table name (not the alias)
        # node.name might be "alias.column" but node.source.this has the real name
        actual_cte_or_table = cte_or_table  # default to parsed name
        if hasattr(node.source, "this") and node.source.this:
            actual_cte_or_table = str(node.source.this)

        # Create namespaced ID using the actual CTE/table name
        if source_type == "Table":
            transform_id = f"table:{actual_cte_or_table}"
            transform_type = "table"
        else:
            transform_id = f"cte:{actual_cte_or_table}"
            transform_type = "cte"

        # Track CTE name -> ID for lookup in second pass (use actual name, not alias)
        cte_to_id[actual_cte_or_table] = transform_id

        # Skip if we've already processed this transformation
        if transform_id in transform_map:
            # But save node data for sources extraction in second pass
            if hasattr(node, "expression") and node.expression:
                nodes_with_data.append((transform_id, node, str(node.expression)))
            continue

        # Build transformation
        transform: dict[str, Any] = {
            "id": transform_id,
            "type": transform_type,
            "column": col_name,
        }

        # Add expression if present and meaningful
        if hasattr(node, "expression") and node.expression:
            expr_sql = str(node.expression)
            if expr_sql and expr_sql.strip() != col_name and len(expr_sql) <= 200:
                transform["expression"] = expr_sql
            elif len(expr_sql) > 200:
                transform["expression"] = expr_sql[:197] + "..."

        transform_map[transform_id] = transform

        # Save for second pass
        if hasattr(node, "expression") and node.expression:
            expr_sql_str = str(node.expression)
            nodes_with_data.append((transform_id, node, expr_sql_str if len(expr_sql_str) <= 200 else expr_sql_str[:200]))

    # Process UNION branches in two passes
    # FIRST: Register all union CTEs in cte_to_id
    for ref_cte, branches in union_branches.items():
        if not branches:
            continue
        transform_id = f"cte:{ref_cte}"
        cte_to_id[ref_cte] = transform_id

    # SECOND: Create union transformations with sources
    for ref_cte, branches in union_branches.items():
        if not branches:
            continue

        # Create union transformation
        transform_id = f"cte:{ref_cte}"
        column = branches[0].get("column", "")  # All branches have same column

        # Build branches array with sources
        formatted_branches: list[dict[str, Any]] = []
        for branch_info in branches:
            branch_entry: dict[str, Any] = {}
            if "expression" in branch_info:
                branch_entry["expression"] = branch_info["expression"]

            # Extract sources for this branch (now all unions are in cte_to_id)
            source_ids: set[str] = set()
            full_expr = branch_info.get("_full_expr", "")
            if full_expr:
                for cte_name in cte_to_id.keys():
                    if f"{cte_name}." in full_expr:
                        source_ids.add(cte_to_id[cte_name])

            branch_entry["sources"] = sorted(list(source_ids))
            formatted_branches.append(branch_entry)

        # Create the union transformation
        union_transform: dict[str, Any] = {
            "id": transform_id,
            "type": "union",
            "column": column,
            "branches": formatted_branches,
        }

        transform_map[transform_id] = union_transform

    # SECOND PASS: Extract sources by looking for CTE references in expressions
    sources_map: dict[str, set[str]] = {}  # transform_id -> set of source ids

    for transform_id, node, expr_str in nodes_with_data:
        source_ids: set[str] = set()

        # If it's a wildcard (SELECT *), get the table from the FROM clause
        if expr_str.strip() == "*" and hasattr(node.source, "find"):
            table_node = node.source.find(exp.Table)
            if table_node and hasattr(table_node, "this"):
                table_name = str(table_node.this)
                source_ids.add(f"table:{table_name}")

        # Look for CTE/table references in the expression (e.g., "orders.order_id")
        for cte_name, cte_id in cte_to_id.items():
            # Don't mark self-reference as a source
            if cte_id != transform_id and f"{cte_name}." in expr_str:
                source_ids.add(cte_id)

        # Also check if expression uses an alias - resolve to actual CTE/table
        # e.g., expression "seg.customer_id" where seg is alias for all_segments
        if "." in expr_str and hasattr(node, "source"):
            # Extract the prefix (potential alias)
            alias_candidate = expr_str.split(".")[0].strip()

            # Try to resolve it from the source
            if hasattr(node.source, "find"):
                table_node = node.source.find(exp.Table)
                if table_node:
                    actual_name = None
                    # Get the actual table/CTE name
                    if hasattr(table_node, "this") and table_node.this:
                        actual_name = str(table_node.this)

                    # Get the alias
                    if hasattr(table_node, "alias_or_name"):
                        alias = str(table_node.alias_or_name)

                        # If expression uses this alias, resolve to actual name
                        # Include alias in output as cte:name[alias] or table:name[alias]
                        if alias == alias_candidate and actual_name and actual_name in cte_to_id:
                            base_id = cte_to_id[actual_name]
                            # Add alias notation if alias differs from actual name
                            if alias != actual_name:
                                source_ids.add(f"{base_id}[{alias}]")
                            else:
                                source_ids.add(base_id)

        # Merge sources for this transformation (handle duplicates from multiple nodes)
        sources_map.setdefault(transform_id, set()).update(source_ids)

    # Build final list with sources
    transformations: list[dict[str, Any]] = []
    for trans in transform_map.values():
        # Union types already have sources in branches, don't add top-level sources
        if trans.get("type") != "union":
            trans["sources"] = sorted(list(sources_map.get(trans["id"], set())))
        transformations.append(trans)

    # Add outer query transformation if we detected it
    if outer_query_sources:
        # Extract column name from any transformation (they all have the same column)
        column_for_query = next((t.get("column") for t in transformations), "")

        # Resolve outer query source references to namespaced IDs
        resolved_sources: list[str] = []
        for ref_name in outer_query_sources:
            # Check if it's a known CTE or table
            if ref_name in cte_to_id:
                resolved_sources.append(cte_to_id[ref_name])
            else:
                # Not in cte_to_id, so it must be a table reference
                # Use the same logic as in FIRST PASS to determine if it's a table
                resolved_sources.append(f"table:{ref_name}")

        outer_query: dict[str, Any] = {
            "id": "query",
            "type": "outer_query",
            "column": column_for_query,
            "sources": sorted(resolved_sources),
        }
        # Insert at the beginning (outer query is first in the chain)
        transformations.insert(0, outer_query)

    return transformations


def _extract_dependencies_from_lineage(
    lineage_node: Any,
    manifest: ManifestLoader | None,
    schema_mapping: dict[str, Any],
    depth: int | None,
    column_name: str = "",
) -> list[dict[str, Any]]:
    """Extract column dependencies from sqlglot lineage node.

    Args:
        lineage_node: Result from sqlglot.lineage()
        manifest: Optional ManifestLoader for dbt resource lookup
        schema_mapping: Schema mapping for table and column resolution
        depth: Maximum depth to traverse
        column_name: Original column being traced (used to replace wildcards)

    Returns:
        List of dependency dicts with column, table, optional dbt_resource,
        and internal CTE transformation path (via_ctes, transformations)
    """
    dependencies: list[dict[str, Any]] = []
    last_cte_column: dict[str, str] = {}  # Track last CTE column before each table

    def walk_dependencies(node: Any, depth_current: int = 0) -> None:
        """Recursively walk lineage tree."""
        if depth is not None and depth_current >= depth:
            return

        for dep in node.walk():
            # Check if this is a table reference (external dependency)
            if hasattr(dep, "source") and hasattr(dep.source, "this"):
                table_name = str(dep.source.this)
                col_name = dep.name

                # Skip our artificial wrapper CTE
                if table_name == "__lineage_final__":
                    continue

                # Track CTE columns (non-Table sources)
                source_type = type(dep.source).__name__
                if source_type != "Table":
                    # This is a CTE/Select - track the column name for the next Table dependency
                    # Extract column name from CTE-qualified references like "orders.order_id"
                    if "." in col_name and col_name != "*":
                        parts = col_name.split(".")
                        cte_col = parts[-1]
                        # Store the most recent non-wildcard column for the next table we'll encounter
                        last_cte_column["__next_table__"] = cte_col

                    # Extract from_table if the Select has a FROM clause with a table
                    if hasattr(dep.source, "find") and callable(dep.source.find):
                        from_table_node = dep.source.find(exp.Table)
                        if from_table_node and hasattr(from_table_node, "this") and from_table_node.this:  # type: ignore[reportAttributeAccessIssue]
                            from_table = str(from_table_node.this)  # type: ignore[reportAttributeAccessIssue]
                            last_cte_column["__next_table_source__"] = from_table

                    continue  # Don't add to dependencies yet

                dependency_info: dict[str, Any] = {
                    "column": col_name,
                    "table": table_name,
                }

                # Always extract db/schema metadata if available
                db = getattr(dep.source, "catalog", None)
                schema_name = getattr(dep.source, "db", None)

                if db:
                    dependency_info["database"] = str(db)
                if schema_name:
                    dependency_info["schema"] = str(schema_name)

                # This is a Table reference - add to dependencies
                # Extract just the column name from potentially qualified references
                # e.g., "customers.first_name" -> "first_name", "id" -> "id"
                final_col_name = col_name.split(".")[-1] if col_name else col_name

                # Replace wildcard with the actual column being traced
                if final_col_name == "*" and column_name:
                    final_col_name = column_name

                # Check if we tracked a CTE column that feeds this table reference
                if "__next_table__" in last_cte_column:
                    source_column = last_cte_column["__next_table__"]
                    # Use the source column as the actual column in this dependency
                    final_col_name = source_column
                    # Clear it after use
                    del last_cte_column["__next_table__"]

                # Update dependency with clean column name
                dependency_info["column"] = final_col_name

                # Try to find the corresponding dbt resource
                if manifest:
                    try:
                        matching_node = manifest.get_resource_node(table_name)
                        if not matching_node.get("multiple_matches"):
                            dependency_info["dbt_resource"] = matching_node.get("unique_id", "")
                    except Exception:
                        pass  # Resource lookup failed, continue with table name

                dependencies.append(dependency_info)

    walk_dependencies(lineage_node)

    # Deduplicate dependencies while preserving order
    seen = set()
    deduplicated = []
    for dep in dependencies:
        # Create a key for deduplication
        key = (dep.get("column"), dep.get("table"), dep.get("database"), dep.get("schema"))
        if key not in seen:
            seen.add(key)
            deduplicated.append(dep)

    return deduplicated


def _extract_cte_path(lineage_node: Any) -> dict[str, Any]:
    """Extract the CTE transformation path from a lineage node.

    Sqlglot lineage nodes have names like:
    - "final.tier" (CTE reference)
    - "aggregated.order_count" (CTE reference)
    - "customer_id" (direct column, no CTE)
    - "base.customer_id" (CTE reference)

    The part before the dot is the CTE name. We walk up the lineage to collect
    all CTEs in the transformation chain.

    Args:
        lineage_node: A single dependency node from sqlglot lineage walk

    Returns:
        Dictionary with:
        - via_ctes: List of CTE names in transformation order (root to leaf)
        - transformations: List of transformation details per step
    """
    via_ctes: list[str] = []
    transformations: list[dict[str, str]] = []
    visited: set[int] = set()  # Track visited nodes by id to prevent cycles

    # Walk up the lineage chain (using downstream references)
    current = lineage_node
    max_iterations = 100  # Hard limit to prevent infinite loops
    iteration = 0

    while current is not None and iteration < max_iterations:
        iteration += 1

        # Prevent cycles by tracking visited nodes
        node_id = id(current)
        if node_id in visited:
            break
        visited.add(node_id)

        # Extract CTE name from node name (format: "cte_name.column_name" or "column_name")
        if hasattr(current, "name") and current.name:
            node_name = current.name

            # Check if this is a CTE reference (has a dot separator)
            if "." in node_name:
                parts = node_name.split(".", 1)
                cte_name = parts[0]
                column_name = parts[1] if len(parts) > 1 else node_name

                # Only add if it's not already in the list (dedup) and looks like a CTE
                # (CTEs typically don't have schema qualifiers like "database.schema.table")
                if cte_name and cte_name not in via_ctes:
                    # Skip if it looks like a database or schema qualifier
                    # (these typically have catalog/db attributes on source)
                    is_table_ref = False
                    if hasattr(current, "source"):
                        source = current.source
                        is_table_ref = hasattr(source, "catalog") or getattr(source, "db", None) is not None

                    if not is_table_ref:
                        via_ctes.append(cte_name)

                        # Capture transformation info
                        transform_info: dict[str, str] = {
                            "cte": cte_name,
                            "column": column_name,
                        }

                        # Try to get expression if available
                        if hasattr(current, "expression") and current.expression is not None:
                            expr_sql = str(current.expression)
                            # Only include if it's not just a simple column reference
                            if expr_sql and expr_sql.strip() != column_name:
                                # Limit expression length to avoid huge outputs
                                if len(expr_sql) > 200:
                                    expr_sql = expr_sql[:197] + "..."
                                transform_info["expression"] = expr_sql

                        transformations.append(transform_info)

        # Move up the lineage chain (downstream attribute points to parent)
        current = getattr(current, "downstream", None) if hasattr(current, "downstream") else None

    return {
        "via_ctes": via_ctes,
        "transformations": transformations,
    }


def _resolve_dependency_resource(
    dependency: dict[str, Any],
    relation_lookup: dict[str, str],
) -> str | None:
    table_name = dependency.get("table")
    if not table_name:
        return None

    database = dependency.get("database")
    schema = dependency.get("schema")

    if database and schema:
        key = _normalize_relation_name(f"{database}.{schema}.{table_name}")
        if key in relation_lookup:
            return relation_lookup[key]

    if schema:
        key = _normalize_relation_name(f"{schema}.{table_name}")
        if key in relation_lookup:
            return relation_lookup[key]

    key = _normalize_relation_name(table_name)
    return relation_lookup.get(key)


def _check_column_in_lineage(lineage_node: Any, source_model: str, source_column: str) -> bool:
    """Check if a source column appears in the lineage tree.

    Args:
        lineage_node: Result from sqlglot.lineage()
        source_model: Model name to look for
        source_column: Column name to look for

    Returns:
        True if the column appears in the lineage
    """
    logger.debug(f"[LINEAGE_CHECK] Looking for {source_model}.{source_column}")

    for dep in lineage_node.walk():
        if hasattr(dep, "name"):
            # dep.name can be either "column_name" or "table.column_name"
            name_parts = dep.name.lower().split(".")
            col_name = name_parts[-1]  # Get the last part (column name)
            table_name = name_parts[0] if len(name_parts) > 1 else None

            logger.debug(f"[LINEAGE_CHECK]   Checking dep.name='{dep.name}' -> col='{col_name}', table='{table_name}'")

            # Check if column matches
            if col_name == source_column.lower():
                logger.debug("[LINEAGE_CHECK]   Column matches! Checking table...")
                # If table name in dep.name, check it matches source_model
                if table_name and source_model.lower() in table_name:
                    logger.debug(f"[LINEAGE_CHECK]   ✓ Table in dep.name matches: {table_name} contains {source_model}")
                    return True
                # If no table name in dep.name, check source attribute
                elif hasattr(dep, "source") and hasattr(dep.source, "this"):
                    source_table = str(dep.source.this).strip('"').lower()
                    logger.debug(f"[LINEAGE_CHECK]   Checking dep.source.this='{source_table}'")
                    if source_model.lower() in source_table:
                        logger.debug(f"[LINEAGE_CHECK]   ✓ Source attribute matches: {source_table} contains {source_model}")
                        return True
                else:
                    logger.debug("[LINEAGE_CHECK]   ✗ Column matches but no table info found")

    logger.debug(f"[LINEAGE_CHECK] ✗ No match found for {source_model}.{source_column}")
    return False


async def _trace_downstream_column(
    manifest: ManifestLoader,
    model_name: str,
    column_name: str,
    depth: int | None,
    dialect: str = "databricks",
    current_depth: int = 0,
) -> list[dict[str, Any]]:
    """Recursively trace where a column is used downstream.

    Args:
        manifest: ManifestLoader instance
        model_name: Model name to start from
        column_name: Column name to trace
        depth: Maximum depth to traverse (None for unlimited)
        dialect: SQL dialect for parsing (default: databricks)
        current_depth: Current recursion depth

    Returns:
        List of downstream usage dictionaries
    """
    logger.info(f"[DOWNSTREAM] Tracing {model_name}.{column_name} at depth {current_depth}")

    if depth is not None and current_depth >= depth:
        return []

    results: list[dict[str, Any]] = []

    # Get downstream models (distance 1)
    try:
        lineage_data = manifest.get_lineage(model_name, resource_type="model", direction="downstream", depth=1)
    except Exception as e:
        logger.warning(f"Could not get downstream lineage for {model_name}: {e}")
        return []

    downstream_models = lineage_data.get("downstream", [])
    logger.info(f"[DOWNSTREAM] Found {len(downstream_models)} downstream models: {[m.get('name') for m in downstream_models]}")

    for downstream_model in downstream_models:
        # Only process models (skip tests, snapshots, etc.)
        if not downstream_model.get("unique_id", "").startswith("model."):
            continue

        # Skip CTE unit test helper models (auto-generated by CTE test generator)
        # Pattern: ends with __<6-char-hash> like customers_enriched__customer_agg__259035
        model_name_check = downstream_model.get("name", "")
        if model_name_check and len(model_name_check) > 8:
            # Check if ends with __<6 hex chars>
            suffix = model_name_check[-8:]
            if suffix[:2] == "__" and all(c in "0123456789abcdef" for c in suffix[2:]):
                logger.debug(f"Skipping CTE test model: {model_name_check}")
                continue

        try:
            # Get downstream model info (schema + SQL) using unified helper
            model_name_downstream = downstream_model["name"]
            downstream_info, compiled_sql, schema_mapping = _prepare_model_analysis(
                manifest,
                model_name_downstream,
            )

            logger.info(f"[DOWNSTREAM] {model_name_downstream}: has_sql={compiled_sql is not None}")

            # Resolve output columns using centralized logic
            output_columns_dict, output_source = _resolve_output_columns(compiled_sql, schema_mapping, downstream_info, dialect)
            if output_columns_dict:
                logger.info(f"[DOWNSTREAM] {model_name_downstream}: using {len(output_columns_dict)} columns from {output_source}")
            else:
                # Fallback: use string search if no column metadata available
                logger.debug(f"[DOWNSTREAM] {model_name_downstream}: No column metadata, using string search")
                source_column_ref_patterns = [
                    f"{model_name}.{column_name}",
                    f".{column_name}",
                    f" {column_name} ",
                    f" {column_name},",
                    f"({column_name}",
                ]
                sql_lower = compiled_sql.lower()
                found_reference = any(pattern.lower() in sql_lower for pattern in source_column_ref_patterns)

                if found_reference:
                    logger.info(f"[DOWNSTREAM] ✓ {model_name_downstream} references {model_name}.{column_name} (string search)")
                    results.append({"model": model_name_downstream, "column": column_name, "distance": current_depth + 1})
                    further_downstream = await _trace_downstream_column(manifest, model_name_downstream, column_name, depth, dialect, current_depth + 1)
                    results.extend(further_downstream)
                continue

            # Check each output column to see if it uses our source column
            logger.debug(f"[DOWNSTREAM] {model_name_downstream}: checking {len(output_columns_dict)} output columns")
            sql_lower = compiled_sql.lower()
            source_column_ref_patterns = [
                f"{model_name}.{column_name}",
                f".{column_name}",
                f" {column_name} ",
                f" {column_name},",
                f"({column_name}",
            ]

            for output_col_name in output_columns_dict.keys():
                try:
                    # Trace this output column's lineage using unified helper
                    column_lineage_result = _analyze_column_lineage(
                        output_col_name,
                        compiled_sql,
                        schema_mapping,
                        model_name_downstream,
                        dialect,
                    )

                    logger.debug(f"[DOWNSTREAM] Check if {output_col_name} uses {model_name}.{column_name}")

                    # Check if our source column appears in the dependencies
                    if _check_column_in_lineage(column_lineage_result, model_name, column_name):
                        logger.info(f"[DOWNSTREAM] ✓ {model_name_downstream}.{output_col_name} USES {model_name}.{column_name}")
                        # This output column uses our source column!
                        results.append({"model": model_name_downstream, "column": output_col_name, "distance": current_depth + 1})

                        # Recurse: trace this column further downstream
                        further_downstream = await _trace_downstream_column(manifest, model_name_downstream, output_col_name, depth, dialect, current_depth + 1)
                        results.extend(further_downstream)
                    else:
                        # Heuristic fallback when lineage doesn't resolve dependencies
                        if output_col_name.lower() == column_name.lower() and any(pattern.lower() in sql_lower for pattern in source_column_ref_patterns):
                            logger.info(f"[DOWNSTREAM] ✓ {model_name_downstream}.{output_col_name} USES {model_name}.{column_name} (heuristic)")
                            results.append({"model": model_name_downstream, "column": output_col_name, "distance": current_depth + 1})

                            further_downstream = await _trace_downstream_column(manifest, model_name_downstream, output_col_name, depth, dialect, current_depth + 1)
                            results.extend(further_downstream)
                        else:
                            logger.debug(f"[DOWNSTREAM] ✗ {model_name_downstream}.{output_col_name} does NOT use {model_name}.{column_name}")

                except SqlglotError as e:
                    # Heuristic fallback when sqlglot fails
                    if output_col_name.lower() == column_name.lower() and any(pattern.lower() in sql_lower for pattern in source_column_ref_patterns):
                        logger.info(f"[DOWNSTREAM] ✓ {model_name_downstream}.{output_col_name} USES {model_name}.{column_name} (heuristic)")
                        results.append({"model": model_name_downstream, "column": output_col_name, "distance": current_depth + 1})

                        further_downstream = await _trace_downstream_column(manifest, model_name_downstream, output_col_name, depth, dialect, current_depth + 1)
                        results.extend(further_downstream)
                    else:
                        logger.warning(f"Could not trace {model_name_downstream}.{output_col_name}: {e}")
                    continue

        except Exception as e:
            # Use get() to avoid UnboundLocalError if exception occurs before model_name_downstream is set
            model_name_for_log = downstream_model.get("name", "unknown")
            logger.warning(f"Error analyzing downstream model {model_name_for_log}: {e}")
            continue

    return results


def _prepare_model_analysis(
    manifest: ManifestLoader,
    model_name: str,
) -> tuple[dict[str, Any], str, dict[str, Any]]:
    """Prepare model for column lineage analysis (unified helper).

    Gets all necessary data in one call:
    - Resource info (metadata, columns)
    - Compiled SQL
    - Schema mapping for sqlglot context

    Args:
        manifest: ManifestLoader instance
        model_name: Model name to analyze

    Returns:
        Tuple of (resource_info, compiled_sql, schema_mapping)

    Raises:
        ValueError: If model not found or has no compiled SQL
    """
    # Get resource info with compiled SQL and database schema
    resource_info = manifest.get_resource_info(
        model_name,
        resource_type="model",
        include_compiled_sql=True,
        include_database_schema=True,
    )

    compiled_sql = resource_info.get("compiled_sql")
    if not compiled_sql:
        raise ValueError(f"No compiled SQL found for model '{model_name}'")

    # Build schema mapping from upstream models
    try:
        upstream_lineage = manifest.get_lineage(
            model_name,
            resource_type="model",
            direction="upstream",
            depth=1,
        )
        schema_mapping = _build_schema_mapping(manifest, upstream_lineage)
    except (ValueError, KeyError, AttributeError) as e:
        logger.warning(f"Could not build schema mapping for {model_name}: {e}")
        schema_mapping = {}

    return resource_info, compiled_sql, schema_mapping


def _clean_static_union_branches(ast: exp.Expression, column_name: str, dialect: str) -> exp.Expression:
    """Remove UNION branches that contain only static literal values.

    When a UNION fails during lineage tracing due to index errors, it's often
    because one branch has all static values (e.g., '-1' as gold_itemkey) while
    another has dynamic column references. This function removes all-static branches
    to allow lineage tracing to proceed.

    Args:
        ast: The SQL AST to clean
        column_name: Column being traced (for logging)
        dialect: SQL dialect

    Returns:
        Cleaned AST with static UNION branches removed
    """
    # Find all UNION nodes (they appear as Union expressions)
    unions_found = list(ast.find_all(exp.Union))

    if not unions_found:
        return ast  # No UNIONs to clean

    logger.debug(f"Found {len(unions_found)} UNION nodes to analyze")

    for union_node in unions_found:
        # UNION in sqlglot is binary: has 'this' (left) and 'expression' (right)
        left_branch = union_node.this
        right_branch = union_node.expression

        # Check if each branch is all-static (contains only literals)
        left_is_static = _is_all_static_branch(left_branch)
        right_is_static = _is_all_static_branch(right_branch)

        logger.debug(f"UNION branch analysis: left_static={left_is_static}, right_static={right_is_static}")

        if left_is_static and not right_is_static:
            # Keep only right branch
            logger.info(f"Removing static left UNION branch for column {column_name}")
            union_node.replace(right_branch)
        elif right_is_static and not left_is_static:
            # Keep only left branch
            logger.info(f"Removing static right UNION branch for column {column_name}")
            union_node.replace(left_branch)
        elif left_is_static and right_is_static:
            # Both static - this is unusual but keep left by convention
            logger.warning(f"Both UNION branches are static for column {column_name}, keeping left")
            union_node.replace(left_branch)
        # else: both dynamic, keep the UNION as-is

    return ast


def _is_all_static_branch(branch: exp.Expression) -> bool:
    """Check if a SELECT branch contains only static literal expressions.

    Args:
        branch: A SELECT expression to analyze

    Returns:
        True if all SELECT expressions are literals, False otherwise
    """
    # Find the SELECT node in this branch
    if isinstance(branch, exp.Select):
        select_node = branch
    else:
        select_node = branch.find(exp.Select)

    if not select_node:
        return False

    # Check all expressions in the SELECT clause
    for select_expr in select_node.expressions:
        # Get the actual expression (might be wrapped in Alias)
        if isinstance(select_expr, exp.Alias):
            actual_expr = select_expr.this
        else:
            actual_expr = select_expr

        # If we find any non-literal, this branch is dynamic
        if not isinstance(actual_expr, exp.Literal):
            # Also check for NULL which is represented differently
            if not (isinstance(actual_expr, exp.Null)):
                return False

    # All expressions are literals
    return True


def _analyze_column_lineage(
    column_name: str,
    compiled_sql: str,
    schema_mapping: dict[str, Any],
    model_name: str,
    dialect: str = "databricks",
) -> Any:
    """Run sqlglot column lineage analysis with automatic UNION cleanup.

    Tries lineage analysis, and if it fails due to UNION structure issues (IndexError),
    automatically cleans static UNION branches and retries. This handles the common
    data warehouse pattern of UNION ALL with hardcoded default rows.

    Args:
        column_name: Column to analyze
        compiled_sql: Compiled SQL
        schema_mapping: Schema context for sqlglot
        model_name: Model name (for error messages)
        dialect: SQL dialect for parsing (default: databricks)

    Returns:
        sqlglot lineage node

    Raises:
        ValueError: If SQL parsing fails or UNION cleanup doesn't resolve issues
    """
    try:
        # Wrap the SQL to enable lineage tracing through SELECT *
        wrapped_ast = _wrap_final_select(compiled_sql, column_name, dialect)

        max_attempts = 3
        for attempt in range(max_attempts):
            try:
                result = lineage(
                    column=column_name,
                    sql=wrapped_ast,
                    schema=schema_mapping,
                    dialect=dialect,
                )

                # Test the result by walking the tree to trigger any IndexError
                for _ in result.walk():
                    pass

                if attempt > 0:
                    logger.info(f"Lineage succeeded after {attempt} UNION cleanup attempt(s) for {model_name}.{column_name}")

                return result  # Success!

            except IndexError as e:
                if "index out of range" in str(e).lower():
                    if attempt == max_attempts - 1:
                        raise ValueError(f"UNION structure issue in {model_name}.{column_name} after {max_attempts} cleanup attempts: {e}")

                    logger.info(f"UNION index error detected for {model_name}.{column_name}, cleaning static branches (attempt {attempt + 1})")
                    # Clean the AST and retry
                    wrapped_ast = _clean_static_union_branches(wrapped_ast, column_name, dialect)
                else:
                    # Different kind of IndexError, re-raise
                    raise

    except SqlglotError as e:
        raise ValueError(f"Failed to parse SQL for column lineage: {e}\nModel: {model_name}, Column: {column_name}")
    except Exception as e:
        logger.exception("Unexpected error in column lineage analysis")
        raise ValueError(f"Column lineage analysis failed: {e}")


def _trace_upstream_recursive(
    manifest: ManifestLoader,
    model_name: str,
    column_name: str,
    depth: int | None,
    dialect: str = "databricks",
    current_depth: int = 0,
    relation_lookup: dict[str, str] | None = None,
    visited: set[str] | None = None,
) -> list[dict[str, Any]]:
    """Recursively trace upstream dependencies for a column.

    Uses unified SQL parsing approach (same as downstream):
    1. Prepare model (get resource_info + compiled_sql + schema_mapping)
    2. Resolve output columns (SQL → warehouse → manifest)
    3. Run lineage per resolved column
    4. Recurse on dependencies

    Args:
        manifest: ManifestLoader instance
        model_name: Model name to analyze
        column_name: Column name to trace
        depth: Maximum depth (None for unlimited)
        dialect: SQL dialect for parsing (default: databricks)
        current_depth: Current recursion depth
        relation_lookup: FQN → unique_id mapping (built once, reused)
        visited: Set of visited unique_ids (prevent cycles)

    Returns:
        List of upstream dependencies with dbt resource mapping
    """
    # Skip wildcards
    if column_name.strip() == "*":
        return []

    # Check depth limit
    if depth is not None and current_depth >= depth:
        return []

    # Initialize tracking on first call
    if relation_lookup is None:
        relation_lookup = _build_relation_lookup(manifest)

    if visited is None:
        visited = set()

    # Prepare model for analysis (unified helper - Phase 1)
    try:
        _, compiled_sql, schema_mapping = _prepare_model_analysis(
            manifest,
            model_name,
        )
    except ValueError as e:
        logger.warning(f"Could not prepare model {model_name}: {e}")
        return []

    # Run lineage analysis (unified helper - Phase 1)
    try:
        column_lineage_result = _analyze_column_lineage(
            column_name,
            compiled_sql,
            schema_mapping,
            model_name,
            dialect,
        )
    except ValueError as e:
        logger.warning(f"Could not analyze column {model_name}.{column_name}: {e}")
        return []

    # Extract root model transformations using new format (id, type, sources)
    root_transformations = _extract_transformations_with_sources(column_lineage_result)

    # Extract dependencies and resolve to dbt resources
    dependencies = _extract_dependencies_from_lineage(
        column_lineage_result,
        manifest,
        schema_mapping,
        depth,
        column_name,
    )

    # Resolve dependency FQNs to dbt resources
    for dependency in dependencies:
        if dependency.get("dbt_resource"):
            continue

        resolved = _resolve_dependency_resource(dependency, relation_lookup)
        if resolved:
            dependency["dbt_resource"] = resolved

    # Enrich model dependencies with their internal transformations
    for dependency in dependencies:
        dbt_resource = dependency.get("dbt_resource")
        if not dbt_resource:
            continue

        node = manifest.get_node_by_unique_id(dbt_resource)
        if not node or node.get("resource_type") != "model":
            continue

        # Get the dependency model's internal transformations
        dep_model_name = node.get("name")
        dep_column = dependency.get("column", "")

        if dep_model_name and dep_column:
            try:
                # Analyze the dependency model to get its internal CTEs
                _, dep_compiled_sql, dep_schema_mapping = _prepare_model_analysis(manifest, dep_model_name)
                dep_lineage = _analyze_column_lineage(dep_column, dep_compiled_sql, dep_schema_mapping, dep_model_name, dialect)

                # Extract internal transformations using new format (id, type, sources)
                internal_transformations = _extract_transformations_with_sources(dep_lineage)

                # Replace the accumulated transformations with internal ones
                if internal_transformations:
                    dependency["transformations"] = internal_transformations
                    # Update via_ctes to match (extract CTE ids)
                    via_ctes = []
                    for t in internal_transformations:
                        if t.get("type") == "cte":
                            # Remove namespace prefix for via_ctes backward compatibility
                            cte_id = t.get("id", "")
                            if cte_id.startswith("cte:"):
                                cte_name = cte_id[4:]  # Strip "cte:" prefix
                                if cte_name not in via_ctes:
                                    via_ctes.append(cte_name)
                    dependency["via_ctes"] = via_ctes

            except Exception as e:
                logger.warning(f"Could not extract internal transformations for {dep_model_name}.{dep_column}: {e}")
                # Keep existing transformations if extraction fails

    # Store results with root model transformations
    results = list(dependencies)

    # Add root transformations metadata to each result (will be extracted by formatter)
    if root_transformations and results:
        for result in results:
            result["__root_transformations__"] = root_transformations
    elif root_transformations:
        # No dependencies found, but we have root transformations - return them anyway
        results = [{"__root_transformations__": root_transformations}]

    # Recurse on each dependency
    for i, dependency in enumerate(dependencies):
        dbt_resource = dependency.get("dbt_resource")
        if not dbt_resource or dbt_resource in visited:
            continue

        node = manifest.get_node_by_unique_id(dbt_resource)
        if not node:
            continue

        if node.get("resource_type") != "model":
            # This is a source/seed - already in results, don't recurse
            continue

        next_model = node.get("name")
        if not next_model:
            continue

        visited.add(dbt_resource)

        # Extract column name from dependency
        dep_column = dependency.get("column", "")

        # Handle wildcards: resolve upstream model's output columns
        if dep_column.endswith(".*") or dep_column.strip() == "*":
            # Use unified helper to prepare upstream model
            try:
                upstream_info, upstream_sql, upstream_schema_mapping = _prepare_model_analysis(
                    manifest,
                    next_model,
                )
            except ValueError:
                continue

            # Resolve output columns using unified logic
            output_columns_dict, _ = _resolve_output_columns(
                upstream_sql,
                upstream_schema_mapping,
                upstream_info,
                dialect,
            )

            # Determine which column to trace through the wildcard
            # Check previous dependency for the actual column name
            if i > 0:
                prev_dep = dependencies[i - 1]
                # Try transformations first, then fall back to column field
                if prev_dep.get("transformations"):
                    column_to_trace = prev_dep["transformations"][0].get("column", column_name)
                else:
                    # Strip CTE prefix if present (e.g., "customers.customer_id" -> "customer_id")
                    prev_column = prev_dep.get("column", "")
                    column_to_trace = prev_column.split(".")[-1] if prev_column else column_name
            else:
                # No previous dependency - use original target column
                column_to_trace = column_name

            # Only trace the specific column we're looking for, not all columns
            if column_to_trace in output_columns_dict:
                results.extend(
                    _trace_upstream_recursive(
                        manifest,
                        next_model,
                        column_to_trace,
                        depth,
                        dialect,
                        current_depth + 1,
                        relation_lookup,
                        visited,
                    )
                )
            continue

        # Regular column: extract name and recurse
        next_column = dep_column.split(".")[-1] if dep_column else column_name

        results.extend(
            _trace_upstream_recursive(
                manifest,
                next_model,
                next_column,
                depth,
                dialect,
                current_depth + 1,
                relation_lookup,
                visited,
            )
        )

    return results

    return results


async def _build_downstream_usages(
    manifest: ManifestLoader,
    model_name: str,
    column_name: str,
    depth: int | None,
    dialect: str,
    current_depth: int = 0,
) -> list[dict[str, Any]]:
    """Build downstream usages with transformations (reversed order).

    For each downstream model using this column, trace how they transform it
    from input to output. Transformations are in reversed order (bottom-to-top:
    table input → CTEs → outer query).

    Args:
        manifest: ManifestLoader instance
        model_name: Source model name
        column_name: Source column name
        depth: Maximum depth to traverse
        dialect: SQL dialect for parsing
        current_depth: Current recursion depth

    Returns:
        List of usage dicts with transformations
    """
    if depth is not None and current_depth >= depth:
        return []

    try:
        lineage_data = manifest.get_lineage(
            model_name,
            resource_type="model",
            direction="downstream",
            depth=1,
        )
    except (ValueError, KeyError, AttributeError) as e:
        logger.warning(f"Could not get downstream lineage for {model_name}: {e}")
        return []

    downstream_models = lineage_data.get("downstream", [])
    usages = []

    for downstream_model in downstream_models:
        if not downstream_model.get("unique_id", "").startswith("model."):
            continue

        # Skip CTE unit test helper models
        model_name_downstream = downstream_model.get("name", "")
        if model_name_downstream and len(model_name_downstream) > 8:
            suffix = model_name_downstream[-8:]
            if suffix[:2] == "__" and all(c in "0123456789abcdef" for c in suffix[2:]):
                continue

        try:
            # Get downstream model's SQL and trace transformations
            downstream_info, compiled_sql, schema_mapping = _prepare_model_analysis(
                manifest,
                model_name_downstream,
            )

            # Resolve output columns to find which column(s) use our source column
            output_columns_dict, _ = _resolve_output_columns(compiled_sql, schema_mapping, downstream_info, dialect)

            # Find output columns that reference our source column
            for output_col_name in output_columns_dict:
                try:
                    # Trace transformations for this output column
                    column_lineage_result = _analyze_column_lineage(
                        column_name=output_col_name,
                        compiled_sql=compiled_sql,
                        schema_mapping=schema_mapping,
                        model_name=model_name_downstream,
                        dialect=dialect,
                    )

                    # Extract transformations (shows how downstream model builds the column)
                    transformations = _extract_transformations_with_sources(column_lineage_result)

                    # Reverse transformations for downstream (bottom-to-top: input → output)
                    # Upstream shows top-to-bottom (query → CTEs → table)
                    # Downstream shows bottom-to-top (table → CTEs → query)
                    transformations = list(reversed(transformations))

                    # Extract dependencies (to find if they reference our source column)
                    deps = _extract_dependencies_from_lineage(
                        column_lineage_result,
                        manifest,
                        schema_mapping,
                        depth=1,  # Only immediate deps
                        column_name=output_col_name,
                    )

                    # Check if any dependency references our source model+column
                    uses_source_column = False
                    for dep in deps:
                        dep_table = dep.get("table", "")
                        dep_column = dep.get("column", "")
                        # Match table name (handle schema-qualified names)
                        if model_name.lower() in dep_table.lower() and dep_column.lower() == column_name.lower():
                            uses_source_column = True
                            break

                    if uses_source_column:
                        # Already extracted transformations above
                        # Recursively get usages of this downstream column
                        nested_usages = await _build_downstream_usages(
                            manifest,
                            model_name_downstream,
                            output_col_name,
                            depth,
                            dialect,
                            current_depth + 1,
                        )

                        usage = {
                            "model": model_name_downstream,
                            "column": output_col_name,
                            "distance": current_depth + 1,
                            "transformations": transformations,
                        }

                        if nested_usages:
                            usage["usages"] = nested_usages

                        usages.append(usage)

                except Exception as e:
                    logger.debug(f"Could not trace {model_name_downstream}.{output_col_name}: {e}")
                    continue

        except Exception as e:
            logger.warning(f"Error analyzing downstream model {model_name_downstream}: {e}")
            continue

    return usages


def _format_lineage_response(model_name: str, column_name: str, direction: str, dependencies: list[dict[str, Any]], downstream_usage: list[dict[str, Any]] | None = None) -> dict[str, Any]:
    """Format the final lineage response.

    Args:
        model_name: Model name
        column_name: Column name
        direction: Direction of lineage
        dependencies: Upstream dependencies
        downstream_usage: Optional downstream usage info (legacy format or new usages format)

    Returns:
        Formatted response dict
    """
    result: dict[str, Any] = {
        "model": model_name,
        "column": column_name,
        "direction": direction,
    }

    # Handle downstream-only format (new unified structure with transformations)
    if direction == "downstream" and downstream_usage is not None:
        # New format: usages with transformations (no root transformations, we're the source)
        result["usages"] = downstream_usage
        return result

    # Extract root transformations from dependencies (if present)
    root_transformations = None
    if dependencies and "__root_transformations__" in dependencies[0]:
        root_transformations = dependencies[0]["__root_transformations__"]
        # Remove from all dependencies (cleanup)
        for dep in dependencies:
            dep.pop("__root_transformations__", None)

    # Add root transformations if present (upstream or both)
    if root_transformations:
        result["transformations"] = root_transformations
        # Extract via_ctes from transformations (new format with id/type)
        via_ctes = []
        for t in root_transformations:
            if t.get("type") == "cte":
                # Extract CTE name from namespaced id (e.g., "cte:orders" -> "orders")
                cte_id = t.get("id", "")
                if cte_id.startswith("cte:"):
                    cte_name = cte_id[4:]  # Strip "cte:" prefix
                    if cte_name not in via_ctes:
                        via_ctes.append(cte_name)
        if via_ctes:
            result["via_ctes"] = via_ctes

    result["dependencies"] = dependencies
    result["dependency_count"] = len(dependencies)

    if downstream_usage is not None and direction != "downstream":
        # For "both" direction: use new usages format
        result["usages"] = downstream_usage

    return result


async def implementation(
    ctx: Context | None,
    model_name: str,
    column_name: str,
    direction: str,
    depth: int | None,
    state: DbtCoreServerContext,
    force_parse: bool = True,
) -> dict[str, Any]:
    """Implementation function for get_column_lineage tool.

    Separated for testing purposes - tests call this directly with explicit state.
    The @tool() decorated get_column_lineage() function calls this with injected dependencies.
    """
    # Initialize state if needed
    await state.ensure_initialized(ctx, force_parse)

    # Verify manifest is available
    if state.manifest is None:
        raise RuntimeError("Manifest not initialized")

    # FIRST: Check if we have compiled code, compile ALL if needed
    # Test source model to see if compilation is needed
    resouce_info = state.manifest.get_resource_info(model_name, resource_type="model", include_compiled_sql=True)
    compiled_sql = resouce_info.get("compiled_sql")
    if not compiled_sql:
        logger.info("No compiled SQL found - compiling entire project")
        runner = await state.get_runner()
        # Compile ALL models (dbt compile with no selector)
        compile_result = await runner.invoke(["compile"])

        if compile_result.success:
            # Reload manifest to get compiled code
            await state.manifest.load()
            # Re-fetch the resource to get updated compiled_code
            resouce_info = state.manifest.get_resource_info(
                model_name,
                resource_type="model",
                include_database_schema=False,
                include_compiled_sql=True,
            )
            compiled_sql = resouce_info.get("compiled_sql")
        else:
            raise RuntimeError(f"Failed to compile project: {compile_result}")

        logger.info("Project compiled successfully")

    # We always need compiled SQL from this point.
    if not compiled_sql:
        raise ValueError(f"No compiled SQL found for model '{model_name}'. Model may not contain SQL code.")

    # Handle multiple matches (check AFTER compilation like get_resource_info does)
    if resouce_info.get("multiple_matches"):
        raise ValueError(f"Multiple models found matching '{model_name}'. Please use unique_id: {[m['unique_id'] for m in resouce_info['matches']]}")

    # Get SQL dialect from manifest metadata
    project_info = state.manifest.get_project_info()
    adapter_type = project_info.get("adapter_type", "databricks")
    dialect = _map_dbt_adapter_to_sqlglot_dialect(adapter_type)
    logger.info(f"Using SQL dialect '{dialect}' for adapter type '{adapter_type}'")

    # Validate that the requested column exists in the model's output
    # Build schema mapping for column resolution
    upstream_lineage = state.manifest.get_lineage(
        model_name,
        resource_type="model",
        direction="upstream",
        depth=1,
    )
    schema_mapping = _build_schema_mapping(state.manifest, upstream_lineage)

    # Resolve output columns
    output_columns, _ = _resolve_output_columns(
        compiled_sql=compiled_sql,
        schema_mapping=schema_mapping,
        resource_info=resouce_info,
        dialect=dialect,
    )

    # Check if requested column exists
    if column_name not in output_columns:
        available_columns = ", ".join(f"'{col}'" for col in sorted(output_columns.keys()))
        raise ValueError(f"Column '{column_name}' not found in output of model '{model_name}'. Available columns: {available_columns}")

    if direction == "downstream":
        # Downstream only - new format with transformations
        usages = await _build_downstream_usages(state.manifest, model_name, column_name, depth, dialect)
        return _format_lineage_response(model_name, column_name, direction, [], usages)

    if direction == "upstream":
        # Upstream only
        dependencies = _trace_upstream_recursive(state.manifest, model_name, column_name, depth, dialect)
        return _format_lineage_response(model_name, column_name, direction, dependencies, None)

    else:  # direction == "both"
        # Both upstream and downstream
        dependencies = _trace_upstream_recursive(state.manifest, model_name, column_name, depth, dialect)
        usages = await _build_downstream_usages(state.manifest, model_name, column_name, depth, dialect)
        return _format_lineage_response(model_name, column_name, direction, dependencies, usages)


@dbtTool()
async def get_column_lineage(
    ctx: Context,
    model_name: str,
    column_name: str,
    direction: str = "upstream",
    depth: int | None = None,
    state: DbtCoreServerContext = Depends(get_state),
) -> dict[str, Any]:
    """Trace column-level lineage through SQL transformations.

    Uses sqlglot to parse compiled SQL and track how columns flow through:
    - CTEs and subqueries
    - JOINs and aggregations
    - Transformations (calculations, CASE statements, etc.)
    - Window functions

    This provides detailed column-to-column dependencies that model-level
    lineage cannot capture.

    Args:
        model_name: Name or unique_id of the dbt model to analyze
        column_name: Name of the column to trace
        direction: Direction to trace lineage:
            - "upstream": Which source columns feed into this column
            - "downstream": Which downstream columns use this column
            - "both": Full bidirectional column lineage
        depth: Maximum levels to traverse (None for unlimited)
            - depth=1: Immediate column dependencies only
            - depth=2: Dependencies + their dependencies
            - None: Full dependency tree

    Returns:
        Column lineage information including:
        - Source columns this column depends on (upstream)
        - Downstream columns that depend on this column
        - Transformations and derivations
        - CTE transformation paths (via_ctes, transformations)
        - dbt resource mapping where available

        Each dependency includes:
        - column: Column name
        - table: Source table name
        - schema: Source schema (if available)
        - database: Source database (if available)
        - via_ctes: List of CTE names in transformation order
        - transformations: Transformation details per CTE step
          - cte: CTE name
          - column: Column name at this step
          - expression: SQL expression (truncated to 200 chars)

    Raises:
        ValueError: If model not found, column not found, or SQL parse fails
        RuntimeError: If sqlglot is not installed

    Examples:
        # Find which source columns feed into revenue
        get_column_lineage("fct_sales", "revenue", "upstream")

        # See what downstream models use customer_id
        get_column_lineage("dim_customers", "customer_id", "downstream")

        # Full bidirectional lineage for a column
        get_column_lineage("fct_orders", "order_total", "both")

    Note:
        Requires sqlglot package. Install with: pip install sqlglot
        The model must be compiled (run 'dbt compile' first).
    """
    return await implementation(ctx, model_name, column_name, direction, depth, state)
