from __future__ import annotations

import typing as t

from sqlglot.optimizer.scope import Scope, build_scope


if t.TYPE_CHECKING:
    from sqlglot._typing import E


def eliminate_ctes(expression: E) -> E:
    """
    Remove unused CTEs from an expression.

    Example:
        >>> import sqlglot
        >>> sql = "WITH y AS (SELECT a FROM x) SELECT a FROM z"
        >>> expression = sqlglot.parse_one(sql)
        >>> eliminate_ctes(expression).sql()
        'SELECT a FROM z'

    Args:
        expression (sqlglot.Expr): expression to optimize
    Returns:
        sqlglot.Expr: optimized expression
    """
    root = build_scope(expression)

    if root:
        ref_count = root.ref_count()

        # Traverse the scope tree in reverse so we can remove chains of unused CTEs
        for scope in reversed(list(root.traverse())):
            if scope.is_cte:
                count = ref_count[id(scope)]
                if count <= 0:
                    cte_node = scope.expression.parent
                    if not cte_node:
                        continue
                    with_node = cte_node.parent
                    cte_node.pop()

                    # Pop the entire WITH clause if this is the last CTE
                    if with_node and len(with_node.expressions) <= 0:
                        with_node.pop()

                    # Decrement the ref count for all sources this CTE selects from
                    for _, source in scope.selected_sources.values():
                        if isinstance(source, Scope):
                            ref_count[id(source)] -= 1

    return expression
