# sql/util.py
# Copyright (C) 2005-2021 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: https://www.opensource.org/licenses/mit-license.php

"""High level utilities which build upon other modules here.

"""

from collections import deque
from itertools import chain

from . import coercions
from . import operators
from . import roles
from . import visitors
from .annotation import _deep_annotate  # noqa
from .annotation import _deep_deannotate  # noqa
from .annotation import _shallow_annotate  # noqa
from .base import _expand_cloned
from .base import _from_objects
from .base import ColumnSet
from .ddl import sort_tables  # noqa
from .elements import _find_columns  # noqa
from .elements import _label_reference
from .elements import _textual_label_reference
from .elements import BindParameter
from .elements import ColumnClause
from .elements import ColumnElement
from .elements import Grouping
from .elements import Label
from .elements import Null
from .elements import UnaryExpression
from .schema import Column
from .selectable import Alias
from .selectable import FromClause
from .selectable import FromGrouping
from .selectable import Join
from .selectable import ScalarSelect
from .selectable import SelectBase
from .selectable import TableClause
from .traversals import HasCacheKey  # noqa
from .. import exc
from .. import util


join_condition = util.langhelpers.public_factory(
    Join._join_condition, ".sql.util.join_condition"
)


def find_join_source(clauses, join_to):
    """Given a list of FROM clauses and a selectable,
    return the first index and element from the list of
    clauses which can be joined against the selectable.  returns
    None, None if no match is found.

    e.g.::

        clause1 = table1.join(table2)
        clause2 = table4.join(table5)

        join_to = table2.join(table3)

        find_join_source([clause1, clause2], join_to) == clause1

    """

    selectables = list(_from_objects(join_to))
    idx = []
    for i, f in enumerate(clauses):
        for s in selectables:
            if f.is_derived_from(s):
                idx.append(i)
    return idx


def find_left_clause_that_matches_given(clauses, join_from):
    """Given a list of FROM clauses and a selectable,
    return the indexes from the list of
    clauses which is derived from the selectable.

    """

    selectables = list(_from_objects(join_from))
    liberal_idx = []
    for i, f in enumerate(clauses):
        for s in selectables:
            # basic check, if f is derived from s.
            # this can be joins containing a table, or an aliased table
            # or select statement matching to a table.  This check
            # will match a table to a selectable that is adapted from
            # that table.  With Query, this suits the case where a join
            # is being made to an adapted entity
            if f.is_derived_from(s):
                liberal_idx.append(i)
                break

    # in an extremely small set of use cases, a join is being made where
    # there are multiple FROM clauses where our target table is represented
    # in more than one, such as embedded or similar.   in this case, do
    # another pass where we try to get a more exact match where we aren't
    # looking at adaption relationships.
    if len(liberal_idx) > 1:
        conservative_idx = []
        for idx in liberal_idx:
            f = clauses[idx]
            for s in selectables:
                if set(surface_selectables(f)).intersection(
                    surface_selectables(s)
                ):
                    conservative_idx.append(idx)
                    break
        if conservative_idx:
            return conservative_idx

    return liberal_idx


def find_left_clause_to_join_from(clauses, join_to, onclause):
    """Given a list of FROM clauses, a selectable,
    and optional ON clause, return a list of integer indexes from the
    clauses list indicating the clauses that can be joined from.

    The presence of an "onclause" indicates that at least one clause can
    definitely be joined from; if the list of clauses is of length one
    and the onclause is given, returns that index.   If the list of clauses
    is more than length one, and the onclause is given, attempts to locate
    which clauses contain the same columns.

    """
    idx = []
    selectables = set(_from_objects(join_to))

    # if we are given more than one target clause to join
    # from, use the onclause to provide a more specific answer.
    # otherwise, don't try to limit, after all, "ON TRUE" is a valid
    # on clause
    if len(clauses) > 1 and onclause is not None:
        resolve_ambiguity = True
        cols_in_onclause = _find_columns(onclause)
    else:
        resolve_ambiguity = False
        cols_in_onclause = None

    for i, f in enumerate(clauses):
        for s in selectables.difference([f]):
            if resolve_ambiguity:
                if set(f.c).union(s.c).issuperset(cols_in_onclause):
                    idx.append(i)
                    break
            elif Join._can_join(f, s) or onclause is not None:
                idx.append(i)
                break

    if len(idx) > 1:
        # this is the same "hide froms" logic from
        # Selectable._get_display_froms
        toremove = set(
            chain(*[_expand_cloned(f._hide_froms) for f in clauses])
        )
        idx = [i for i in idx if clauses[i] not in toremove]

    # onclause was given and none of them resolved, so assume
    # all indexes can match
    if not idx and onclause is not None:
        return range(len(clauses))
    else:
        return idx


def visit_binary_product(fn, expr):
    """Produce a traversal of the given expression, delivering
    column comparisons to the given function.

    The function is of the form::

        def my_fn(binary, left, right)

    For each binary expression located which has a
    comparison operator, the product of "left" and
    "right" will be delivered to that function,
    in terms of that binary.

    Hence an expression like::

        and_(
            (a + b) == q + func.sum(e + f),
            j == r
        )

    would have the traversal::

        a <eq> q
        a <eq> e
        a <eq> f
        b <eq> q
        b <eq> e
        b <eq> f
        j <eq> r

    That is, every combination of "left" and
    "right" that doesn't further contain
    a binary comparison is passed as pairs.

    """
    stack = []

    def visit(element):
        if isinstance(element, ScalarSelect):
            # we don't want to dig into correlated subqueries,
            # those are just column elements by themselves
            yield element
        elif element.__visit_name__ == "binary" and operators.is_comparison(
            element.operator
        ):
            stack.insert(0, element)
            for l in visit(element.left):
                for r in visit(element.right):
                    fn(stack[0], l, r)
            stack.pop(0)
            for elem in element.get_children():
                visit(elem)
        else:
            if isinstance(element, ColumnClause):
                yield element
            for elem in element.get_children():
                for e in visit(elem):
                    yield e

    list(visit(expr))
    visit = None  # remove gc cycles


def find_tables(
    clause,
    check_columns=False,
    include_aliases=False,
    include_joins=False,
    include_selects=False,
    include_crud=False,
):
    """locate Table objects within the given expression."""

    tables = []
    _visitors = {}

    if include_selects:
        _visitors["select"] = _visitors["compound_select"] = tables.append

    if include_joins:
        _visitors["join"] = tables.append

    if include_aliases:
        _visitors["alias"] = _visitors["subquery"] = _visitors[
            "tablesample"
        ] = _visitors["lateral"] = tables.append

    if include_crud:
        _visitors["insert"] = _visitors["update"] = _visitors[
            "delete"
        ] = lambda ent: tables.append(ent.table)

    if check_columns:

        def visit_column(column):