from __future__ import annotations

from dataclasses import replace
import re

from .ir import (
    SQL2BinaryExpr,
    SQL2CastExpr,
    SQL2ColumnExpr,
    SQL2CreateViewOp,
    SQL2ExistsExpr,
    SQL2Expr,
    SQL2FromItem,
    SQL2FunctionExpr,
    SQL2Join,
    SQL2LiteralExpr,
    SQL2ProgramIR,
    SQL2Query,
    SQL2RecursiveQuery,
    SQL2SemiNaiveLoop,
    SQL2SemiNaiveTable,
    SQL2SelectItem,
    SQL2SelectQuery,
    SQL2StratumProgram,
    SQL2SubquerySource,
    SQL2TableSource,
    SQL2UnionQuery,
    SQL2WindowExpr,
    SQL2OrderByExpr,
)
from .type_inference import is_text_sql_type, normalize_sql_type, same_sql_type


class SQL2ProgramOptimizer:
    """IR-to-IR optimizer for SQL2 program structure."""

    def optimize_query(self, query: SQL2Query) -> SQL2Query:
        optimized_query, _ = self._rewrite_query(query)
        return optimized_query

    def optimize(self, program: SQL2ProgramIR) -> SQL2ProgramIR:
        externally_referenced: set[str] = set()
        optimized_reversed: list[SQL2StratumProgram] = []
        for stratum in reversed(program.strata):
            optimized = self._optimize_stratum(stratum, externally_referenced=externally_referenced)
            optimized_reversed.append(optimized)
            externally_referenced.update(self._stratum_table_dependencies(optimized))
        return SQL2ProgramIR(strata=tuple(reversed(optimized_reversed)))

    def _optimize_stratum(self, stratum: SQL2StratumProgram, *, externally_referenced: set[str]) -> SQL2StratumProgram:
        operations = list(stratum.operations)
        while True:
            changed = False

            operations, rewritten = self._rewrite_operation_queries(operations)
            changed = changed or rewritten

            pruned = self._prune_dead_internal_views(operations, externally_referenced=externally_referenced)
            if len(pruned) != len(operations):
                operations = pruned
                changed = True

            operations, inlined = self._inline_single_ref_internal_view(
                operations,
                externally_referenced=externally_referenced,
            )
            changed = changed or inlined

            if not changed:
                break
        loops: list[SQL2SemiNaiveLoop] = []
        for loop in stratum.semi_naive_loops:
            tables: list[SQL2SemiNaiveTable] = []
            for table in loop.tables:
                base_query, _ = self._rewrite_query(table.base_query)
                delta_query, _ = self._rewrite_query(table.delta_query)
                merge_query, _ = self._rewrite_query(table.merge_query)
                final_query, _ = self._rewrite_query(table.final_query)
                tables.append(
                    SQL2SemiNaiveTable(
                        table_name=table.table_name,
                        columns=table.columns,
                        base_query=base_query,
                        delta_query=delta_query,
                        merge_query=merge_query,
                        final_query=final_query,
                        accum_view_name=table.accum_view_name,
                        delta_view_name=table.delta_view_name,
                        next_delta_view_name=table.next_delta_view_name,
                    )
                )
            loops.append(SQL2SemiNaiveLoop(tables=tuple(tables)))
        return SQL2StratumProgram(
            stratum_id=stratum.stratum_id,
            operations=tuple(operations),
            semi_naive_loops=tuple(loops),
        )

    def _stratum_table_dependencies(self, stratum: SQL2StratumProgram) -> set[str]:
        deps: set[str] = set()
        for op in stratum.operations:
            deps.update(self._query_table_dependencies(op.query))
        for loop in stratum.semi_naive_loops:
            for table in loop.tables:
                deps.update(self._query_table_dependencies(table.base_query))
                deps.update(self._query_table_dependencies(table.delta_query))
                deps.update(self._query_table_dependencies(table.merge_query))
                deps.update(self._query_table_dependencies(table.final_query))
        return deps

    def _rewrite_operation_queries(self, operations: list[SQL2CreateViewOp]) -> tuple[list[SQL2CreateViewOp], bool]:
        changed = False
        rewritten_ops: list[SQL2CreateViewOp] = []
        for op in operations:
            rewritten_query, query_changed = self._rewrite_query(op.query)
            changed = changed or query_changed
            rewritten_ops.append(SQL2CreateViewOp(name=op.name, query=rewritten_query))
        return rewritten_ops, changed

    def _rewrite_query(self, query: SQL2Query) -> tuple[SQL2Query, bool]:
        if isinstance(query, SQL2UnionQuery):
            changed = False
            parts: list[SQL2Query] = []
            for part in query.parts:
                rewritten_part, part_changed = self._rewrite_query(part)
                changed = changed or part_changed
                parts.append(rewritten_part)
            if len(parts) == 1:
                return parts[0], True
            rewritten_union = SQL2UnionQuery(parts=tuple(parts), distinct=query.distinct)
            return rewritten_union, changed
        if isinstance(query, SQL2RecursiveQuery):
            changed = False
            base_query = query.base_query
            if base_query is not None:
                rewritten_base, base_changed = self._rewrite_query(base_query)
                base_query = rewritten_base
                changed = changed or base_changed
            recursive_query = query.recursive_query
            if recursive_query is not None:
                rewritten_recursive, recursive_changed = self._rewrite_query(recursive_query)
                recursive_query = rewritten_recursive
                changed = changed or recursive_changed
            final_query = query.final_query
            if final_query is not None:
                rewritten_final, final_changed = self._rewrite_query(final_query)
                if isinstance(rewritten_final, SQL2SelectQuery):
                    final_query = rewritten_final
                changed = changed or final_changed
            rewritten_recursive_query = SQL2RecursiveQuery(
                cte_name=query.cte_name,
                column_names=query.column_names,
                base_query=base_query,
                recursive_query=recursive_query,
                final_query=final_query,
                distinct=query.distinct,
            )
            return rewritten_recursive_query, changed

        changed = False
        from_item: SQL2FromItem | None = None
        if query.from_item is not None:
            from_item, from_changed = self._rewrite_from_item(query.from_item)
            changed = changed or from_changed
        joins: list[SQL2Join] = []
        for join in query.joins:
            rewritten_join, join_changed = self._rewrite_join(join)
            changed = changed or join_changed
            joins.append(rewritten_join)
        where: list[SQL2Expr] = []
        for expr in query.where:
            rewritten_expr, expr_changed = self._rewrite_expr(expr)
            changed = changed or expr_changed
            where.append(rewritten_expr)
        group_by: list[SQL2Expr] = []
        for expr in query.group_by:
            rewritten_expr, expr_changed = self._rewrite_expr(expr)
            changed = changed or expr_changed
            group_by.append(rewritten_expr)
        having: list[SQL2Expr] = []
        for expr in query.having:
            rewritten_expr, expr_changed = self._rewrite_expr(expr)
            changed = changed or expr_changed
            having.append(rewritten_expr)
        select_items: list[SQL2SelectItem] = []
        for item in query.select_items:
            rewritten_expr, expr_changed = self._rewrite_expr(item.expr)
            changed = changed or expr_changed
            select_items.append(SQL2SelectItem(expr=rewritten_expr, alias=item.alias))
        deduped_where, deduped_where_changed = self._dedupe_conjunctive_predicates(tuple(where))
        deduped_having, deduped_having_changed = self._dedupe_conjunctive_predicates(tuple(having))
        changed = changed or deduped_where_changed or deduped_having_changed

        rewritten_select = SQL2SelectQuery(
            select_items=tuple(select_items),
            from_item=from_item,
            joins=tuple(joins),
            where=deduped_where,
            group_by=tuple(group_by),
            having=deduped_having,
            distinct=query.distinct,
        )

        while True:
            local_changed = False
            flattened = self._flatten_nested_merge_query(rewritten_select)
            if flattened != rewritten_select:
                rewritten_select = flattened
                local_changed = True
            unwrapped = self._unwrap_passthrough_wrapper(rewritten_select)
            if unwrapped != rewritten_select:
                rewritten_select = unwrapped
                local_changed = True
            changed = changed or local_changed
            if not local_changed:
                break
        return rewritten_select, changed

    def _rewrite_from_item(self, from_item: SQL2FromItem) -> tuple[SQL2FromItem, bool]:
        source = from_item.source
        if isinstance(source, SQL2SubquerySource):
            rewritten_query, changed = self._rewrite_query(source.query)
            return (
                SQL2FromItem(
                    source=SQL2SubquerySource(query=rewritten_query),
                    alias=from_item.alias,
                    lateral=from_item.lateral,
                ),
                changed,
            )
        return from_item, False

    def _rewrite_join(self, join: SQL2Join) -> tuple[SQL2Join, bool]:
        source, source_changed = self._rewrite_from_item(join.source)
        changed = source_changed
        on_exprs: list[SQL2Expr] = []
        for expr in join.on:
            rewritten_expr, expr_changed = self._rewrite_expr(expr)
            changed = changed or expr_changed
            on_exprs.append(rewritten_expr)
        deduped_on, deduped_on_changed = self._dedupe_conjunctive_predicates(tuple(on_exprs))
        changed = changed or deduped_on_changed
        return SQL2Join(source=source, join_type=join.join_type, on=deduped_on), changed

    def _rewrite_expr(self, expr: SQL2Expr) -> tuple[SQL2Expr, bool]:
        if isinstance(expr, SQL2ExistsExpr):
            rewritten_query, changed = self._rewrite_query(expr.query)
            rewritten_exists = SQL2ExistsExpr(query=rewritten_query, negated=expr.negated)
            simplified = self._simplify_exists_where_only(rewritten_exists)
            if isinstance(simplified, SQL2ExistsExpr):
                simplified = self._simplify_nested_exists(simplified)
            return simplified, changed or (simplified != rewritten_exists)
        if isinstance(expr, SQL2BinaryExpr):
            left, left_changed = self._rewrite_expr(expr.left)
            right, right_changed = self._rewrite_expr(expr.right)
            return SQL2BinaryExpr(op=expr.op, left=left, right=right), left_changed or right_changed
        if isinstance(expr, SQL2FunctionExpr):
            changed = False
            args: list[SQL2Expr] = []
            for arg in expr.args:
                rewritten_arg, arg_changed = self._rewrite_expr(arg)
                changed = changed or arg_changed
                args.append(rewritten_arg)
            return SQL2FunctionExpr(name=expr.name, args=tuple(args), distinct=expr.distinct), changed
        if isinstance(expr, SQL2WindowExpr):
            changed = False
            args: list[SQL2Expr] = []
            for arg in expr.args:
                rewritten_arg, arg_changed = self._rewrite_expr(arg)
                changed = changed or arg_changed
                args.append(rewritten_arg)
            partition_by: list[SQL2Expr] = []
            for part in expr.partition_by:
                rewritten_part, part_changed = self._rewrite_expr(part)
                changed = changed or part_changed
                partition_by.append(rewritten_part)
            order_by: list[SQL2OrderByExpr] = []
            for item in expr.order_by:
                rewritten_expr, item_changed = self._rewrite_expr(item.expr)
                changed = changed or item_changed
                order_by.append(SQL2OrderByExpr(expr=rewritten_expr, descending=item.descending))
            return (
                SQL2WindowExpr(
                    name=expr.name,
                    args=tuple(args),
                    partition_by=tuple(partition_by),
                    order_by=tuple(order_by),
                    distinct=expr.distinct,
                ),
                changed,
            )
        if isinstance(expr, SQL2CastExpr):
            rewritten_expr, changed = self._rewrite_expr(expr.expr)
            if isinstance(rewritten_expr, SQL2CastExpr) and self._same_sql_type(
                rewritten_expr.target_type, expr.target_type
            ):
                return rewritten_expr, True
            if (
                isinstance(rewritten_expr, SQL2LiteralExpr)
                and isinstance(rewritten_expr.value, str)
                and is_text_sql_type(expr.target_type)
            ):
                return rewritten_expr, True
            return SQL2CastExpr(expr=rewritten_expr, target_type=expr.target_type), changed
        return expr, False

    def _simplify_nested_exists(self, expr: SQL2ExistsExpr) -> SQL2Expr:
        query = expr.query
        if not isinstance(query, SQL2SelectQuery):
            return expr
        if query.from_item is not None or query.joins or query.group_by or query.having or query.distinct:
            return expr
        if len(query.where) != 1:
            return expr
        nested = query.where[0]
        if not isinstance(nested, SQL2ExistsExpr):
            return expr
        return SQL2ExistsExpr(query=nested.query, negated=(expr.negated ^ nested.negated))

    def _simplify_exists_where_only(self, expr: SQL2ExistsExpr) -> SQL2Expr:
        query = expr.query
        if not isinstance(query, SQL2SelectQuery):
            return expr
        if query.from_item is not None or query.joins or query.group_by or query.having or query.distinct:
            return expr
        if not query.where:
            return SQL2LiteralExpr(value=not expr.negated)
        predicate = query.where[0]
        for extra in query.where[1:]:
            predicate = SQL2BinaryExpr(op="AND", left=predicate, right=extra)
        normalized = SQL2FunctionExpr(
            name="COALESCE",
            args=(predicate, SQL2LiteralExpr(value=False)),
        )
        if expr.negated:
            return SQL2BinaryExpr(op="=", left=normalized, right=SQL2LiteralExpr(value=False))
        return normalized

    def _dedupe_conjunctive_predicates(
        self,
        predicates: tuple[SQL2Expr, ...],
    ) -> tuple[tuple[SQL2Expr, ...], bool]:
        seen: set[tuple[object, ...]] = set()
        deduped: list[SQL2Expr] = []
        changed = False
        for predicate in predicates:
            for conjunct in self._flatten_conjunctive_exprs(predicate):
                key = self._predicate_dedupe_key(conjunct)
                if key in seen:
                    changed = True
                    continue
                seen.add(key)
                deduped.append(conjunct)
        if len(deduped) != len(predicates):
            changed = True
        return tuple(deduped), changed

    def _flatten_conjunctive_exprs(self, predicate: SQL2Expr) -> tuple[SQL2Expr, ...]:
        if isinstance(predicate, SQL2BinaryExpr) and predicate.op.upper() == "AND":
            left = self._flatten_conjunctive_exprs(predicate.left)
            right = self._flatten_conjunctive_exprs(predicate.right)
            return (*left, *right)
        return (predicate,)

    def _predicate_dedupe_key(self, predicate: SQL2Expr) -> tuple[object, ...]:
        if isinstance(predicate, SQL2BinaryExpr):
            op = predicate.op.upper()
            if op in {"=", "==", "!="}:
                left_key = self._expr_dedupe_repr(predicate.left)
                right_key = self._expr_dedupe_repr(predicate.right)
                if right_key < left_key:
                    left_key, right_key = right_key, left_key
                return ("binary_commutative", op, left_key, right_key)
        return ("expr", self._expr_dedupe_repr(predicate))

    def _expr_dedupe_repr(self, expr: SQL2Expr) -> str:
        return repr(expr)

    def _normalize_sql_type(self, sql_type: str) -> str:
        return normalize_sql_type(sql_type)

    def _same_sql_type(self, left: str, right: str) -> bool:
        return same_sql_type(left, right)

    def _unwrap_passthrough_wrapper(self, query: SQL2SelectQuery) -> SQL2SelectQuery:
        if query.distinct or query.joins or query.where or query.group_by or query.having:
            return query
        from_item = query.from_item
        if from_item is None or not isinstance(from_item.source, SQL2SubquerySource):
            return query
        inner = from_item.source.query
        if not isinstance(inner, SQL2SelectQuery):
            return query
        projected_columns: list[str] = []
        for item in query.select_items:
            if not isinstance(item.expr, SQL2ColumnExpr):
                return query
            if item.expr.table_alias != from_item.alias:
                return query
            if item.alias is None or item.alias != item.expr.column_name:
                return query
            projected_columns.append(item.expr.column_name)
        inner_aliases = [item.alias for item in inner.select_items]
        if len(projected_columns) != len(inner_aliases):
            return query
        if any(alias is None for alias in inner_aliases):
            return query
        if tuple(projected_columns) != tuple(inner_aliases):
            return query
        return inner

    def _flatten_nested_merge_query(self, query: SQL2SelectQuery) -> SQL2SelectQuery:
        merge_sig = self._merge_query_signature(query)
        if merge_sig is None:
            return query
        from_item = query.from_item
        assert from_item is not None
        assert isinstance(from_item.source, SQL2SubquerySource)
        union_query = from_item.source.query
        assert isinstance(union_query, SQL2UnionQuery)

        flattened = False
        flattened_parts: list[SQL2Query] = []
        for part in union_query.parts:
            if isinstance(part, SQL2SelectQuery):
                part_sig = self._merge_query_signature(part)
                if part_sig == merge_sig:
                    part_from = part.from_item
                    assert part_from is not None
                    assert isinstance(part_from.source, SQL2SubquerySource)
                    part_union = part_from.source.query
                    assert isinstance(part_union, SQL2UnionQuery)
                    flattened_parts.extend(part_union.parts)
                    flattened = True
                    continue
            flattened_parts.append(part)
        if not flattened:
            return query
        return SQL2SelectQuery(
            select_items=query.select_items,
            from_item=SQL2FromItem(
                source=SQL2SubquerySource(
                    query=SQL2UnionQuery(parts=tuple(flattened_parts), distinct=union_query.distinct)
                ),
                alias=from_item.alias,
                lateral=from_item.lateral,
            ),
            joins=query.joins,
            where=query.where,
            group_by=query.group_by,
            having=query.having,
            distinct=query.distinct,
        )

    def _merge_query_signature(self, query: SQL2SelectQuery) -> tuple[tuple[str, ...], tuple[str, ...]] | None:
        if query.joins or query.where or query.having or query.distinct:
            return None
        from_item = query.from_item
        if from_item is None or not isinstance(from_item.source, SQL2SubquerySource):
            return None
        if not isinstance(from_item.source.query, SQL2UnionQuery):
            return None
        if not query.group_by:
            return None
        group_alias = None
        key_columns: list[str] = []
        for expr in query.group_by:
            if not isinstance(expr, SQL2ColumnExpr):
                return None
            if group_alias is None:
                group_alias = expr.table_alias
            elif group_alias != expr.table_alias:
                return None
            key_columns.append(expr.column_name)
        if group_alias is None:
            return None

        key_set = set(key_columns)
        seen_keys: set[str] = set()
        value_columns: list[str] = []
        for item in query.select_items:
            if item.alias is None:
                return None
            expr = item.expr
            if isinstance(expr, SQL2ColumnExpr):
                if expr.table_alias != group_alias or expr.column_name != item.alias:
                    return None
                if expr.column_name not in key_set:
                    return None
                seen_keys.add(expr.column_name)
                continue
            if isinstance(expr, SQL2FunctionExpr) and expr.name.upper() == "MAX" and len(expr.args) == 1:
                max_arg = expr.args[0]
                if not isinstance(max_arg, SQL2ColumnExpr):
                    return None
                if max_arg.table_alias != group_alias or max_arg.column_name != item.alias:
                    return None
                if item.alias in key_set:
                    return None
                value_columns.append(item.alias)
                continue
            return None
        if seen_keys != key_set:
            return None
        return tuple(key_columns), tuple(value_columns)

    def _prune_dead_internal_views(
        self,
        operations: list[SQL2CreateViewOp],
        *,
        externally_referenced: set[str],
    ) -> list[SQL2CreateViewOp]:
        if not operations:
            return []
        ops_by_name = {op.name: op for op in operations}
        deps_by_name = {
            op.name: {dep for dep in self._query_table_dependencies(op.query) if dep in ops_by_name}
            for op in operations
        }
        roots = [
            op.name
            for op in operations
            if (not self._is_internal_view_name(op.name)) or (op.name in externally_referenced)
        ]
        if not roots:
            # Keep all operations if no explicit public roots are present.
            return list(operations)

        keep: set[str] = set()
        stack = list(roots)
        while stack:
            name = stack.pop()
            if name in keep:
                continue
            keep.add(name)
            stack.extend(dep for dep in deps_by_name.get(name, ()) if dep not in keep)
        return [op for op in operations if op.name in keep]

    def _is_internal_view_name(self, name: str) -> bool:
        last = name.rsplit(".", 1)[-1]
        return bool(re.match(r"^sql2_s\d+_.*$", last))

    def _inline_single_ref_internal_view(
        self,
        operations: list[SQL2CreateViewOp],
        *,
        externally_referenced: set[str],
    ) -> tuple[list[SQL2CreateViewOp], bool]:
        if len(operations) < 2:
            return operations, False
        ops_by_name = {op.name: op for op in operations}
        op_order = {op.name: ix for ix, op in enumerate(operations)}

        ref_counts: dict[str, int] = {name: 0 for name in ops_by_name}
        referrers: dict[str, set[str]] = {name: set() for name in ops_by_name}
        for op in operations:
            counts = self._query_table_ref_counts(op.query)
            for name, count in counts.items():
                if name not in ref_counts:
                    continue
                ref_counts[name] += count
                referrers[name].add(op.name)

        candidates = [
            name
            for name in ops_by_name
            if self._is_internal_view_name(name)
            and name not in externally_referenced
            and ref_counts.get(name, 0) == 1
            and len(referrers.get(name, set())) == 1
        ]
        if not candidates:
            return operations, False
        candidates.sort(key=lambda name: op_order[name])

        for name in candidates:
            referrer_names = referrers[name]
            if not referrer_names:
                continue
            referrer_name = next(iter(referrer_names))
            if referrer_name == name:
                # Do not inline self-recursive/internal self-referencing views.
                continue
            referrer_op = ops_by_name.get(referrer_name)
            target_op = ops_by_name.get(name)
            if referrer_op is None or target_op is None:
                continue
            # Inlining `name` into `referrer` would create a cycle if the
            # inlined query depends on `referrer` itself (for example deferred
            # state updates that seed from the prior public table).
            if referrer_name in self._query_table_dependencies(target_op.query):
                continue
            replaced_query = self._replace_table_with_subquery(
                referrer_op.query,
                table_name=name,
                replacement=target_op.query,
            )
            updated_referrer = SQL2CreateViewOp(name=referrer_op.name, query=replaced_query)
            next_operations: list[SQL2CreateViewOp] = []
            for op in operations:
                if op.name == name:
                    continue
                if op.name == referrer_name:
                    next_operations.append(updated_referrer)
                else:
                    next_operations.append(op)
            return next_operations, True
        return operations, False

    def _replace_table_with_subquery(self, query: SQL2Query, *, table_name: str, replacement: SQL2Query) -> SQL2Query:
        if isinstance(query, SQL2UnionQuery):
            parts = tuple(self._replace_table_with_subquery(part, table_name=table_name, replacement=replacement) for part in query.parts)
            return replace(query, parts=parts)
        if isinstance(query, SQL2RecursiveQuery):
            base_query = (
                self._replace_table_with_subquery(query.base_query, table_name=table_name, replacement=replacement)
                if query.base_query is not None
                else None
            )
            recursive_query = (
                self._replace_table_with_subquery(query.recursive_query, table_name=table_name, replacement=replacement)
                if query.recursive_query is not None
                else None
            )
            final_query = (
                self._replace_table_with_subquery(query.final_query, table_name=table_name, replacement=replacement)
                if query.final_query is not None
                else None
            )
            return SQL2RecursiveQuery(
                cte_name=query.cte_name,
                column_names=query.column_names,
                base_query=base_query,
                recursive_query=recursive_query,
                final_query=final_query if isinstance(final_query, SQL2SelectQuery) else query.final_query,
                distinct=query.distinct,
            )
        from_item = self._replace_from_item_source(query.from_item, table_name=table_name, replacement=replacement) if query.from_item else None
        joins = tuple(self._replace_join_source(join, table_name=table_name, replacement=replacement) for join in query.joins)
        where = tuple(self._replace_expr_tables(expr, table_name=table_name, replacement=replacement) for expr in query.where)
        having = tuple(self._replace_expr_tables(expr, table_name=table_name, replacement=replacement) for expr in query.having)
        select_items = tuple(
            SQL2SelectItem(
                expr=self._replace_expr_tables(item.expr, table_name=table_name, replacement=replacement),
                alias=item.alias,
            )
            for item in query.select_items
        )
        return SQL2SelectQuery(
            select_items=select_items,
            from_item=from_item,
            joins=joins,
            where=where,
            group_by=query.group_by,
            having=having,
            distinct=query.distinct,
        )

    def _replace_join_source(self, join: SQL2Join, *, table_name: str, replacement: SQL2Query) -> SQL2Join:
        source = self._replace_from_item_source(join.source, table_name=table_name, replacement=replacement)
        on = tuple(self._replace_expr_tables(expr, table_name=table_name, replacement=replacement) for expr in join.on)
        return SQL2Join(source=source, join_type=join.join_type, on=on)

    def _replace_from_item_source(self, from_item: SQL2FromItem, *, table_name: str, replacement: SQL2Query) -> SQL2FromItem:
        source = from_item.source
        if isinstance(source, SQL2TableSource):
            if source.name == table_name:
                return SQL2FromItem(
                    source=SQL2SubquerySource(query=replacement),
                    alias=from_item.alias,
                    lateral=from_item.lateral,
                )
            return from_item
        if isinstance(source, SQL2SubquerySource):
            rewritten = self._replace_table_with_subquery(source.query, table_name=table_name, replacement=replacement)
            return SQL2FromItem(
                source=SQL2SubquerySource(query=rewritten),
                alias=from_item.alias,
                lateral=from_item.lateral,
            )
        return from_item

    def _replace_expr_tables(self, expr: SQL2Expr, *, table_name: str, replacement: SQL2Query) -> SQL2Expr:
        if isinstance(expr, SQL2ExistsExpr):
            rewritten = self._replace_table_with_subquery(expr.query, table_name=table_name, replacement=replacement)
            return SQL2ExistsExpr(query=rewritten, negated=expr.negated)
        if isinstance(expr, SQL2BinaryExpr):
            left = self._replace_expr_tables(expr.left, table_name=table_name, replacement=replacement)
            right = self._replace_expr_tables(expr.right, table_name=table_name, replacement=replacement)
            return SQL2BinaryExpr(op=expr.op, left=left, right=right)
        if isinstance(expr, SQL2FunctionExpr):
            args = tuple(self._replace_expr_tables(arg, table_name=table_name, replacement=replacement) for arg in expr.args)
            return SQL2FunctionExpr(name=expr.name, args=args, distinct=expr.distinct)
        if isinstance(expr, SQL2WindowExpr):
            args = tuple(self._replace_expr_tables(arg, table_name=table_name, replacement=replacement) for arg in expr.args)
            partition_by = tuple(
                self._replace_expr_tables(arg, table_name=table_name, replacement=replacement)
                for arg in expr.partition_by
            )
            order_by = tuple(
                SQL2OrderByExpr(
                    expr=self._replace_expr_tables(item.expr, table_name=table_name, replacement=replacement),
                    descending=item.descending,
                )
                for item in expr.order_by
            )
            return SQL2WindowExpr(
                name=expr.name,
                args=args,
                partition_by=partition_by,
                order_by=order_by,
                distinct=expr.distinct,
            )
        if isinstance(expr, SQL2CastExpr):
            rewritten = self._replace_expr_tables(expr.expr, table_name=table_name, replacement=replacement)
            return SQL2CastExpr(expr=rewritten, target_type=expr.target_type)
        return expr

    def _query_table_dependencies(self, query: SQL2Query) -> set[str]:
        return set(self._query_table_ref_counts(query))

    def _query_table_ref_counts(self, query: SQL2Query) -> dict[str, int]:
        if isinstance(query, SQL2UnionQuery):
            deps: dict[str, int] = {}
            for part in query.parts:
                self._merge_ref_counts(deps, self._query_table_ref_counts(part))
            return deps
        if isinstance(query, SQL2RecursiveQuery):
            deps: dict[str, int] = {}
            if query.base_query is not None:
                self._merge_ref_counts(deps, self._query_table_ref_counts(query.base_query))
            if query.recursive_query is not None:
                self._merge_ref_counts(deps, self._query_table_ref_counts(query.recursive_query))
            if query.final_query is not None:
                self._merge_ref_counts(deps, self._query_table_ref_counts(query.final_query))
            deps.pop(query.cte_name, None)
            return deps
        deps: dict[str, int] = {}
        if query.from_item is not None:
            self._merge_ref_counts(deps, self._from_item_table_dependencies(query.from_item))
        for join in query.joins:
            self._merge_ref_counts(deps, self._from_item_table_dependencies(join.source))
            for predicate in join.on:
                self._merge_ref_counts(deps, self._expr_query_dependencies(predicate))
        for predicate in query.where:
            self._merge_ref_counts(deps, self._expr_query_dependencies(predicate))
        for predicate in query.having:
            self._merge_ref_counts(deps, self._expr_query_dependencies(predicate))
        for item in query.select_items:
            self._merge_ref_counts(deps, self._expr_query_dependencies(item.expr))
        return deps

    def _from_item_table_dependencies(self, from_item: SQL2FromItem) -> dict[str, int]:
        source = from_item.source
        if isinstance(source, SQL2TableSource):
            return {source.name: 1}
        if isinstance(source, SQL2SubquerySource):
            return self._query_table_ref_counts(source.query)
        return {}

    def _expr_query_dependencies(self, expr: SQL2Expr) -> dict[str, int]:
        if isinstance(expr, SQL2ExistsExpr):
            return self._query_table_ref_counts(expr.query)
        if isinstance(expr, SQL2BinaryExpr):
            deps = self._expr_query_dependencies(expr.left)
            self._merge_ref_counts(deps, self._expr_query_dependencies(expr.right))
            return deps
        if isinstance(expr, SQL2FunctionExpr):
            deps: dict[str, int] = {}
            for arg in expr.args:
                self._merge_ref_counts(deps, self._expr_query_dependencies(arg))
            return deps
        if isinstance(expr, SQL2WindowExpr):
            deps: dict[str, int] = {}
            for arg in expr.args:
                self._merge_ref_counts(deps, self._expr_query_dependencies(arg))
            for part in expr.partition_by:
                self._merge_ref_counts(deps, self._expr_query_dependencies(part))
            for order_item in expr.order_by:
                self._merge_ref_counts(deps, self._expr_query_dependencies(order_item.expr))
            return deps
        if isinstance(expr, SQL2CastExpr):
            return self._expr_query_dependencies(expr.expr)
        return {}

    def _merge_ref_counts(self, target: dict[str, int], source: dict[str, int]) -> None:
        for name, count in source.items():
            target[name] = target.get(name, 0) + count
