from __future__ import annotations

import typing as t

from collections import defaultdict

from sqlglot import expressions as exp
from sqlglot.helper import find_new_name, seq_get
from sqlglot.optimizer.scope import Scope, traverse_scope

if t.TYPE_CHECKING:
    from sqlglot._typing import E

    FromOrJoin = t.Union[exp.From, exp.Join]


def merge_subqueries(expression: E, leave_tables_isolated: bool = False) -> E:
    """
    Rewrite sqlglot AST to merge derived tables into the outer query.

    This also merges CTEs if they are selected from only once.

    Example:
        >>> import sqlglot
        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
        >>> merge_subqueries(expression).sql()
        'SELECT x.a FROM x CROSS JOIN y'

    If `leave_tables_isolated` is True, this will not merge inner queries into outer
    queries if it would result in multiple table selects in a single query:
        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
        >>> merge_subqueries(expression, leave_tables_isolated=True).sql()
        'SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y'

    Inspired by https://dev.mysql.com/doc/refman/8.0/en/derived-table-optimization.html

    Args:
        expression (sqlglot.Expr): expression to optimize
        leave_tables_isolated (bool):
    Returns:
        sqlglot.Expr: optimized expression
    """
    # Shared across both passes so the scope tree is only built once, as long as merge_ctes
    # doesn't mutate the AST; if it does, the scopes it was given are no longer valid, so the
    # scope tree needs to be rebuilt before merge_derived_tables runs.
    scopes = traverse_scope(expression)
    expression, merged_ctes = merge_ctes(expression, leave_tables_isolated, scopes=scopes)

    if merged_ctes:
        scopes = traverse_scope(expression)

    expression = merge_derived_tables(expression, leave_tables_isolated, scopes=scopes)
    return expression


# If a derived table has these Select args, it can't be merged
UNMERGABLE_ARGS = set(exp.Select.arg_types) - {
    "expressions",
    "from_",
    "joins",
    "where",
    "order",
    "hint",
}


# Projections in the outer query that are instances of these types can be replaced
# without getting wrapped in parentheses, because the precedence won't be altered.
SAFE_TO_REPLACE_UNWRAPPED = (
    exp.Column,
    exp.EQ,
    exp.Func,
    exp.NEQ,
    exp.Paren,
)


def merge_ctes(
    expression: E,
    leave_tables_isolated: bool = False,
    scopes: list[Scope] | None = None,
) -> tuple[E, bool]:
    # All places where we select from CTEs.
    # We key on the CTE scope so we can detect CTES that are selected from multiple times.
    cte_selections = defaultdict(list)
    for outer_scope in traverse_scope(expression) if scopes is None else scopes:
        for table, inner_scope in outer_scope.selected_sources.values():
            if isinstance(inner_scope, Scope) and inner_scope.is_cte:
                cte_selections[id(inner_scope)].append(
                    (
                        outer_scope,
                        inner_scope,
                        table,
                    )
                )

    merged = False
    singular_cte_selections = [v[0] for k, v in cte_selections.items() if len(v) == 1]
    for outer_scope, inner_scope, table in singular_cte_selections:
        from_or_join = table.find_ancestor(exp.From, exp.Join)
        if not isinstance(from_or_join, (exp.From, exp.Join)):
            continue
        if _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
            alias = table.alias_or_name
            _rename_inner_sources(outer_scope, inner_scope, alias)
            _merge_from(outer_scope, inner_scope, table, alias)
            _merge_expressions(outer_scope, inner_scope, alias)
            _merge_order(outer_scope, inner_scope)
            _merge_joins(outer_scope, inner_scope, from_or_join)
            _merge_where(outer_scope, inner_scope, from_or_join)
            _merge_hints(outer_scope, inner_scope)
            _pop_cte(inner_scope)
            outer_scope.clear_cache()
            merged = True
    return expression, merged


def merge_derived_tables(
    expression: E,
    leave_tables_isolated: bool = False,
    scopes: list[Scope] | None = None,
) -> E:
    for outer_scope in traverse_scope(expression) if scopes is None else scopes:
        for subquery in outer_scope.derived_tables:
            from_or_join = subquery.find_ancestor(exp.From, exp.Join)
            if not isinstance(from_or_join, (exp.From, exp.Join)):
                continue
            alias = subquery.alias_or_name
            inner_scope = outer_scope.sources[alias]
            if not isinstance(inner_scope, Scope):
                continue
            if _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
                _rename_inner_sources(outer_scope, inner_scope, alias)
                _merge_from(outer_scope, inner_scope, subquery, alias)
                _merge_expressions(outer_scope, inner_scope, alias)
                _merge_order(outer_scope, inner_scope)
                _merge_joins(outer_scope, inner_scope, from_or_join)
                _merge_where(outer_scope, inner_scope, from_or_join)
                _merge_hints(outer_scope, inner_scope)
                outer_scope.clear_cache()

    return expression


def _mergeable(
    outer_scope: Scope,
    inner_scope: Scope,
    leave_tables_isolated: bool,
    from_or_join: FromOrJoin,
) -> bool:
    """
    Return True if `inner_select` can be merged into outer query.
    """
    inner_select = inner_scope.expression.unnest()
    outer = outer_scope.expression
    outer_args = outer.args
    inner_name = from_or_join.alias_or_name

    def _window_projection_blocks_merge(window_aliases: set[str]) -> bool:
        """A window function's result depends on the full row set it sees, so merging the
        subquery into the outer query is unsafe when:
          - the outer query filters or joins (WHERE/JOIN), which changes that row set, or
          - a window column is referenced in an operation that isn't pushed down
            (GROUP BY, ORDER BY, HAVING, aggregate).
        """
        if not window_aliases:
            return False

        if outer_args.get("where") or outer_args.get("joins"):
            return True

        return any(
            column.table == inner_name
            and column.name in window_aliases
            and column.find_ancestor(exp.Group, exp.Order, exp.Having, exp.AggFunc)
            for column in outer_scope.columns
        )

    def _literal_group_unmergeable(number_literal_aliases: set[str]) -> bool:
        """
        A numeric-literal projection referenced in GROUP BY can't be inlined, because a bare
        integer literal is positional. A reference that is itself a top-level GROUP BY item
        can merge as the ordinal of the outer projection that selects it; any other reference,
        e.g., ROLLUP / CUBE / GROUPING SETS, tuples, expressions, etc, blocks the merge, since
        ordinals aren't universally supported there, e.g., Presto / Trino only accept columns.
        """
        group = outer_args.get("group")
        if not group:
            return False

        if not number_literal_aliases:
            return False

        grouped = set()
        top_level_ids = {id(e.unnest()) for e in group.expressions}
        for col in group.find_all(exp.Column):
            if col.table != inner_name or col.name not in number_literal_aliases:
                continue
            if id(col) not in top_level_ids:
                return True
            grouped.add(col.name)

        if not grouped:
            return False

        projected = set()
        for s in t.cast(exp.Select, outer).selects:
            unaliased = s.unalias()
            if isinstance(unaliased, exp.Column) and unaliased.table == inner_name:
                projected.add(unaliased.name)

        return not grouped <= projected

    def _outer_select_joins_on_inner_select_join(projections: dict[str, exp.Expr]) -> bool:
        """
        All columns from the inner select in the ON clause must be from the first FROM table.

        That is, this can be merged:
            SELECT * FROM x JOIN (SELECT y.a AS a FROM y JOIN z) AS q ON x.a = q.a
                                         ^^^           ^
        But this can't:
            SELECT * FROM x JOIN (SELECT z.a AS a FROM y JOIN z) AS q ON x.a = q.a
                                         ^^^                  ^
        """
        if not isinstance(from_or_join, exp.Join):
            return False

        on = from_or_join.args.get("on")
        if not on:
            return False
        selections = [c.name for c in on.find_all(exp.Column) if c.table == inner_name]
        inner_from = inner_scope.expression.args.get("from_")
        if not inner_from:
            return False
        inner_from_table = inner_from.alias_or_name
        return any(
            col.table != inner_from_table
            for selection in selections
            for col in projections[selection].find_all(exp.Column)
        )

    def _is_recursive() -> bool:
        # Recursive CTEs look like this:
        #     WITH RECURSIVE cte AS (
        #       SELECT * FROM x  <-- inner scope
        #       UNION ALL
        #       SELECT * FROM cte  <-- outer scope
        #     )
        cte = inner_scope.expression.parent
        node = outer.parent

        while node:
            if node is cte:
                return True
            node = node.parent
        return False

    def _literal_in_order_by(number_literal_aliases: set[str]) -> bool:
        """A numeric-literal projection under a bare ORDER BY key can't merge (would become positional)."""
        order = outer_args.get("order")
        if not order:
            return False
        ordered = {
            key.name
            for o in order.expressions
            if isinstance(key := o.this.unnest(), exp.Column) and key.table == inner_name
        }
        return bool(number_literal_aliases & ordered)

    if (
        not isinstance(outer, exp.Select)
        or outer.is_star
        or not isinstance(inner_select, exp.Select)
        or any(v for k, v in inner_select.args.items() if k in UNMERGABLE_ARGS)
        or inner_select.args.get("from_") is None
        or outer_scope.pivots
        or (leave_tables_isolated and len(outer_scope.selected_sources) > 1)
        or (isinstance(from_or_join, exp.Join) and inner_select.args.get("joins"))
        or (
            isinstance(from_or_join, exp.Join)
            and inner_select.args.get("where")
            and from_or_join.side in ("FULL", "LEFT", "RIGHT")
        )
        or (
            isinstance(from_or_join, exp.From)
            and inner_select.args.get("where")
            and any(j.side in ("FULL", "RIGHT") for j in outer_args.get("joins", []))
        )
        or (inner_select.args.get("order") and outer_scope.is_union)
        or isinstance(seq_get(inner_select.expressions, 0), exp.QueryTransform)
    ):
        return False

    # Single pass over the projections: replaces the separate AggFunc/Select/Explode and
    # Window tree walks, and precomputes what the checks below need per-projection.
    window_aliases: set[str] = set()
    number_literal_aliases: set[str] = set()
    projections: dict[str, exp.Expr] = {}

    for s in inner_select.selects:
        name = s.alias_or_name
        projections[name] = s
        if s.unalias().is_number:
            number_literal_aliases.add(name)
        for node in s.walk():
            if isinstance(node, (exp.AggFunc, exp.Select, exp.Explode)):
                return False
            if isinstance(node, exp.Window):
                window_aliases.add(name)

    return (
        not _outer_select_joins_on_inner_select_join(projections)
        and not _window_projection_blocks_merge(window_aliases)
        and not _literal_group_unmergeable(number_literal_aliases)
        and not _literal_in_order_by(number_literal_aliases)
        and not (inner_scope.is_cte and _is_recursive())
    )


def _rename_inner_sources(outer_scope: Scope, inner_scope: Scope, alias: str) -> None:
    """
    Renames any sources in the inner query that conflict with names in the outer query.
    """
    inner_taken = set(inner_scope.selected_sources)
    outer_taken = set(outer_scope.selected_sources)
    conflicts = outer_taken.intersection(inner_taken)
    conflicts -= {alias}

    taken = outer_taken.union(inner_taken)

    for conflict in conflicts:
        new_name = find_new_name(taken, conflict)

        source, _ = inner_scope.selected_sources[conflict]
        new_alias = exp.to_identifier(new_name)

        if isinstance(source, exp.Table) and source.alias:
            source.set("alias", exp.TableAlias(this=new_alias))
        elif isinstance(source, exp.Table):
            source.replace(exp.alias_(source, new_alias))
        elif isinstance(source.parent, exp.Subquery):
            source.parent.set("alias", exp.TableAlias(this=new_alias))

        for column in inner_scope.source_columns(conflict):
            column.set("table", exp.to_identifier(new_name))

        inner_scope.rename_source(conflict, new_name)


def _merge_from(
    outer_scope: Scope,
    inner_scope: Scope,
    node_to_replace: exp.Selectable,
    alias: str,
) -> None:
    """
    Merge FROM clause of inner query into outer query.
    """
    new_subquery = inner_scope.expression.args["from_"].this
    new_subquery.set("joins", node_to_replace.args.get("joins"))
    node_to_replace.replace(new_subquery)
    for join_hint in outer_scope.join_hints:
        tables = join_hint.find_all(exp.Table)
        for table in tables:
            if table.alias_or_name == node_to_replace.alias_or_name:
                table.set("this", exp.to_identifier(new_subquery.alias_or_name))
    outer_scope.remove_source(alias)
    outer_scope.add_source(
        new_subquery.alias_or_name, inner_scope.sources[new_subquery.alias_or_name]
    )


def _merge_joins(outer_scope: Scope, inner_scope: Scope, from_or_join: FromOrJoin) -> None:
    """
    Merge JOIN clauses of inner query into outer query.
    """

    new_joins = []

    joins = inner_scope.expression.args.get("joins") or []

    for join in joins:
        new_joins.append(join)
        outer_scope.add_source(join.alias_or_name, inner_scope.sources[join.alias_or_name])

    if new_joins:
        outer_joins = outer_scope.expression.args.get("joins", [])

        # Maintain the join order
        if isinstance(from_or_join, exp.From):
            position = 0
        else:
            position = outer_joins.index(from_or_join) + 1
        outer_joins[position:position] = new_joins

        outer_scope.expression.set("joins", outer_joins)


def _merge_expressions(outer_scope: Scope, inner_scope: Scope, alias: str) -> None:
    """
    Merge projections of inner query into outer query.

    Args:
        outer_scope (sqlglot.optimizer.scope.Scope)
        inner_scope (sqlglot.optimizer.scope.Scope)
        alias (str)
    """
    # Collect all columns that reference the alias of the inner query
    outer_columns = defaultdict(list)
    for column in outer_scope.columns:
        if column.table == alias:
            outer_columns[column.name].append(column)

    group = outer_scope.expression.args.get("group")

    # Replace columns with the projection expression in the inner query
    for expression in inner_scope.expression.expressions:
        projection_name = expression.alias_or_name
        if not projection_name:
            continue
        columns_to_replace = outer_columns.get(projection_name, [])
        if not columns_to_replace:
            continue

        expression = expression.unalias()
        must_wrap_expression = not isinstance(expression, SAFE_TO_REPLACE_UNWRAPPED)

        is_number = expression.is_number
        last = len(columns_to_replace) - 1

        group_ordinal = None
        if is_number and isinstance(outer_scope.expression, exp.Select):
            # Find the ordinal of the outer SELECT that references this inner projection.
            for j, s in enumerate(outer_scope.expression.selects):
                unaliased = s.unalias()
                if (
                    isinstance(unaliased, exp.Column)
                    and unaliased.table == alias
                    and unaliased.name == projection_name
                ):
                    group_ordinal = j + 1
                    break

        for i, column in enumerate(columns_to_replace):
            parent = column.parent

            # A numeric-literal projection can't be inlined into a top-level GROUP BY item
            # (positional context), canonicalize to the projection's ordinal to match
            # qualify. _mergeable guarantees the ordinal exists.
            if is_number and group:
                item = next((e for e in group.expressions if e.unnest() is column), None)
                if item is not None:
                    item.replace(exp.Literal.number(group_ordinal))
                    continue

            # Ensures we don't alter the intended operator precedence if there's additional
            # context surrounding the outer expression (i.e. it's not a simple projection).
            if isinstance(parent, (exp.Unary, exp.Binary)) and must_wrap_expression:
                expression = exp.paren(expression, copy=False)

            # make sure we do not accidentally change the name of the column
            if isinstance(parent, exp.Select) and column.name != expression.name:
                expression = exp.alias_(expression, column.name)

            # Skip the expensive deep copy for the last reference since the inner query
            # is about to be removed, so we can move the expression directly
            column.replace(expression.copy() if i < last else expression)


def _merge_where(outer_scope: Scope, inner_scope: Scope, from_or_join: FromOrJoin) -> None:
    """
    Merge WHERE clause of inner query into outer query.

    Args:
        outer_scope (sqlglot.optimizer.scope.Scope)
        inner_scope (sqlglot.optimizer.scope.Scope)
        from_or_join (exp.From|exp.Join)
    """
    where = inner_scope.expression.args.get("where")
    if not where or not where.this:
        return

    expression = outer_scope.expression

    if isinstance(from_or_join, exp.Join):
        # Merge predicates from an outer join to the ON clause
        # if it only has columns that are already joined
        from_ = expression.args.get("from_")
        sources = {from_.alias_or_name} if from_ else set()

        for join in expression.args["joins"]:
            source = join.alias_or_name
            sources.add(source)
            if source == from_or_join.alias_or_name:
                break

        if exp.column_table_names(where.this) <= sources:
            from_or_join.on(where.this, copy=False)
            from_or_join.set("on", from_or_join.args.get("on"))
            return

    t.cast(exp.Select, expression).where(where.this, copy=False)


def _merge_order(outer_scope: Scope, inner_scope: Scope) -> None:
    """
    Merge ORDER clause of inner query into outer query.

    Args:
        outer_scope (sqlglot.optimizer.scope.Scope)
        inner_scope (sqlglot.optimizer.scope.Scope)
    """
    inner_order = inner_scope.expression.args.get("order")
    if not inner_order:
        return

    if (
        any(
            outer_scope.expression.args.get(arg) for arg in ["group", "distinct", "having", "order"]
        )
        or len(outer_scope.selected_sources) != 1
        or any(expression.find(exp.AggFunc) for expression in outer_scope.expression.expressions)
    ):
        return

    outer_scope.expression.set("order", inner_order)


def _merge_hints(outer_scope: Scope, inner_scope: Scope) -> None:
    inner_scope_hint = inner_scope.expression.args.get("hint")
    if not inner_scope_hint:
        return
    outer_scope_hint = outer_scope.expression.args.get("hint")
    if outer_scope_hint:
        for hint_expression in inner_scope_hint.expressions:
            outer_scope_hint.append("expressions", hint_expression)
    else:
        outer_scope.expression.set("hint", inner_scope_hint)


def _pop_cte(inner_scope: Scope) -> None:
    """
    Remove CTE from the AST.

    Args:
        inner_scope (sqlglot.optimizer.scope.Scope)
    """
    cte = inner_scope.expression.parent
    if not cte:
        return
    with_ = cte.parent
    if not with_:
        return
    if len(with_.expressions) == 1:
        with_.pop()
    else:
        cte.pop()
