# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
# mypy: no-warn-return-any, allow-any-generics

from __future__ import annotations

from collections.abc import Iterable
from collections.abc import Iterator
import contextlib
from itertools import chain
import re
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Protocol
from typing import TYPE_CHECKING
from typing import TypeGuard
from typing import TypeVar
from typing import Union

from sqlalchemy import __version__
from sqlalchemy import exc as sa_exc
from sqlalchemy import schema
from sqlalchemy import sql
from sqlalchemy import types as sqltypes
from sqlalchemy.schema import CheckConstraint
from sqlalchemy.schema import Column
from sqlalchemy.schema import ForeignKeyConstraint
from sqlalchemy.schema import Table
from sqlalchemy.sql import visitors
from sqlalchemy.sql.base import _NoneName
from sqlalchemy.sql.base import DialectKWArgs
from sqlalchemy.sql.elements import BindParameter
from sqlalchemy.sql.elements import ColumnClause
from sqlalchemy.sql.elements import conv
from sqlalchemy.sql.elements import TextClause
from sqlalchemy.sql.elements import UnaryExpression
from sqlalchemy.sql.naming import _constraint_name_for_table
from sqlalchemy.sql.naming import _NONE_NAME as _NONE_NAME  # type: ignore[attr-defined] # noqa: E501
from sqlalchemy.sql.visitors import traverse

from .messaging import warn

if TYPE_CHECKING:
    from sqlalchemy import ClauseElement
    from sqlalchemy import Identity
    from sqlalchemy import Index
    from sqlalchemy.engine import Connection
    from sqlalchemy.engine import Dialect
    from sqlalchemy.engine import Transaction
    from sqlalchemy.sql.base import ColumnCollection
    from sqlalchemy.sql.compiler import SQLCompiler
    from sqlalchemy.sql.elements import ColumnElement
    from sqlalchemy.sql.schema import Constraint
    from sqlalchemy.sql.schema import ForeignKey
    from sqlalchemy.sql.schema import SchemaItem

_CE = TypeVar("_CE", bound=Union["ColumnElement[Any]", "SchemaItem"])


class _CompilerProtocol(Protocol):
    def __call__(self, element: Any, compiler: Any, **kw: Any) -> str: ...


def _safe_int(value: str) -> int | str:
    try:
        return int(value)
    except:
        return value


_vers = tuple(
    [_safe_int(x) for x in re.findall(r"(\d+|[abc]\d)", __version__)]
)
sqla_2_0_25 = _vers >= (2, 25)
sqla_2_1 = _vers >= (2, 1)
sqlalchemy_version = __version__

if TYPE_CHECKING:

    def compiles(
        element: type[ClauseElement], *dialects: str
    ) -> Callable[[_CompilerProtocol], _CompilerProtocol]: ...

else:
    from sqlalchemy.ext.compiler import compiles  # noqa: I100,I202


identity_has_dialect_kwargs = issubclass(schema.Identity, DialectKWArgs)


def _get_identity_options_dict(
    identity: Identity | schema.Sequence | None,
    dialect_kwargs: bool = False,
) -> dict[str, Any]:
    if identity is None:
        return {}
    elif identity_has_dialect_kwargs:
        assert hasattr(identity, "_as_dict")
        as_dict = identity._as_dict()
        if dialect_kwargs:
            assert isinstance(identity, DialectKWArgs)
            as_dict.update(identity.dialect_kwargs)
    else:
        as_dict = {}
        if isinstance(identity, schema.Identity):
            # always=None means something different than always=False
            as_dict["always"] = identity.always
            if identity.on_null is not None:
                as_dict["on_null"] = identity.on_null
        # attributes common to Identity and Sequence
        attrs = (
            "start",
            "increment",
            "minvalue",
            "maxvalue",
            "nominvalue",
            "nomaxvalue",
            "cycle",
            "cache",
            "order",
        )
        as_dict.update(
            {
                key: getattr(identity, key, None)
                for key in attrs
                if getattr(identity, key, None) is not None
            }
        )
    return as_dict


_ConstraintName = Union[None, str, _NoneName]
_ConstraintNameDefined = Union[str, _NoneName]


def constraint_name_defined(
    name: _ConstraintName,
) -> TypeGuard[_ConstraintNameDefined]:
    return name is _NONE_NAME or isinstance(name, (str, _NoneName))


def constraint_name_string(name: _ConstraintName) -> TypeGuard[str]:
    return isinstance(name, str)


def constraint_name_or_none(name: _ConstraintName) -> str | None:
    return name if constraint_name_string(name) else None


AUTOINCREMENT_DEFAULT = "auto"


@contextlib.contextmanager
def _ensure_scope_for_ddl(
    connection: Connection | None,
) -> Iterator[None]:
    try:
        in_transaction = connection.in_transaction  # type: ignore[union-attr]
    except AttributeError:
        # catch for MockConnection, None
        in_transaction = None
        pass

    # yield outside the catch
    if in_transaction is None:
        yield
    else:
        if not in_transaction():
            assert connection is not None
            with connection.begin():
                yield
        else:
            yield


def _safe_begin_connection_transaction(
    connection: Connection,
) -> Transaction:
    transaction = connection.get_transaction()
    if transaction:
        return transaction
    else:
        return connection.begin()


def _safe_commit_connection_transaction(
    connection: Connection,
) -> None:
    transaction = connection.get_transaction()
    if transaction:
        transaction.commit()


def _safe_rollback_connection_transaction(
    connection: Connection,
) -> None:
    transaction = connection.get_transaction()
    if transaction:
        transaction.rollback()


def _get_connection_in_transaction(connection: Connection | None) -> bool:
    try:
        in_transaction = connection.in_transaction  # type: ignore
    except AttributeError:
        # catch for MockConnection
        return False
    else:
        return in_transaction()


def _idx_table_bound_expressions(idx: Index) -> Iterable[ColumnElement[Any]]:
    return idx.expressions  # type: ignore


def _copy(schema_item: _CE, **kw) -> _CE:
    if hasattr(schema_item, "_copy"):
        return schema_item._copy(**kw)
    else:
        return schema_item.copy(**kw)  # type: ignore[union-attr]


def _connectable_has_table(
    connectable: Connection, tablename: str, schemaname: str | None
) -> bool:
    return connectable.dialect.has_table(connectable, tablename, schemaname)


def _exec_on_inspector(inspector, statement, **params):
    with inspector._operation_context() as conn:
        return conn.execute(statement, params)


def _nullability_might_be_unset(metadata_column):
    from sqlalchemy.sql import schema

    return metadata_column._user_defined_nullable is schema.NULL_UNSPECIFIED


def _server_default_is_computed(*server_default) -> bool:
    return any(isinstance(sd, schema.Computed) for sd in server_default)


def _server_default_is_identity(*server_default) -> bool:
    return any(isinstance(sd, schema.Identity) for sd in server_default)


def _table_for_constraint(constraint: Constraint) -> Table:
    if isinstance(constraint, ForeignKeyConstraint):
        table = constraint.parent
        assert table is not None
        return table  # type: ignore[return-value]
    elif isinstance(constraint, CheckConstraint):
        parent = constraint.parent
        if not isinstance(parent, Table):
            table = parent.table
            assert table is not None
            return table
        else:
            return parent
    else:
        return constraint.table


def _columns_for_constraint(constraint):
    if isinstance(constraint, ForeignKeyConstraint):
        return [fk.parent for fk in constraint.elements]
    elif isinstance(constraint, CheckConstraint):
        return _find_columns(constraint.sqltext)
    else:
        return list(constraint.columns)


def _resolve_for_variant(type_, dialect):
    if _type_has_variants(type_):
        base_type, mapping = _get_variant_mapping(type_)
        return mapping.get(dialect.name, base_type)
    else:
        return type_


if hasattr(sqltypes.TypeEngine, "_variant_mapping"):  # 2.0

    def _type_has_variants(type_):
        return bool(type_._variant_mapping)

    def _get_variant_mapping(type_):
        return type_, type_._variant_mapping

else:

    def _type_has_variants(type_):
        return type(type_) is sqltypes.Variant

    def _get_variant_mapping(type_):
        return type_.impl, type_.mapping


def _get_table_key(name: str, schema: str | None) -> str:
    if schema is None:
        return name
    else:
        return schema + "." + name


def _fk_spec(constraint: ForeignKeyConstraint) -> Any:
    if TYPE_CHECKING:
        assert constraint.columns is not None
        assert constraint.elements is not None
        assert isinstance(constraint.parent, Table)

    source_columns = [
        constraint.columns[key].name for key in constraint.column_keys
    ]

    source_table = constraint.parent.name
    source_schema = constraint.parent.schema
    target_schema = constraint.elements[0].column.table.schema
    target_table = constraint.elements[0].column.table.name
    target_columns = [element.column.name for element in constraint.elements]
    ondelete = constraint.ondelete
    onupdate = constraint.onupdate
    deferrable = constraint.deferrable
    initially = constraint.initially
    return (
        source_schema,
        source_table,
        source_columns,
        target_schema,
        target_table,
        target_columns,
        onupdate,
        ondelete,
        deferrable,
        initially,
    )


class _ForeignKeyTarget(NamedTuple):
    """Stand-in for SQLAlchemy 2.1's ``ForeignKeyTarget`` named tuple."""

    schema: str | None
    table_name: str
    column_name: str | None


def _fk_target_tokens(fk: ForeignKey) -> _ForeignKeyTarget:
    """Return the ``(schema, table_name, column_name)`` names that identify
    the target of a :class:`.ForeignKey`.

    These names are the representation SQLAlchemy computes against; unlike
    the dotted string form of ``ForeignKey._get_colspec()`` they remain
    unambiguous when a name itself contains a dot.

    """
    fk_any: Any = fk
    if sqla_2_1:
        return _ForeignKeyTarget(*fk_any.target_tokens)
    else:
        return _ForeignKeyTarget(*fk_any._column_tokens)


def _fk_target_table_key(fk: ForeignKey) -> str | None:
    """Return the :attr:`.MetaData.tables` key under which the target of a
    :class:`.ForeignKey` is, or would be, registered.

    Returns ``None`` when the target was given as a :class:`.Column` that is
    not associated with a :class:`.Table`, in which case no key can be named.

    """
    fk_any: Any = fk
    if sqla_2_1:
        return fk_any.target_table_key
    else:
        return fk_any._table_key()


def _fk_target_is_named(fk: ForeignKey) -> bool:
    """Return True if a :class:`.ForeignKey` identifies its target by name
    rather than by a :class:`.Column` object.

    """
    fk_any: Any = fk
    if sqla_2_1:
        return fk_any.target_column is None
    else:
        return isinstance(fk_any._colspec, str)


def _fk_is_self_referential(constraint: ForeignKeyConstraint) -> bool:
    assert constraint.parent is not None
    return (
        _fk_target_table_key(constraint.elements[0]) == constraint.parent.key
    )


def _is_type_bound(constraint: Constraint) -> bool:
    # this deals with SQLAlchemy #3260, don't copy CHECK constraints
    # that will be generated by the type.
    # new feature added for #3260
    return constraint._type_bound


def _conv_constraint_names(table: Table) -> None:
    """Mark the names of constraints and indexes on ``table`` as final.

    A name that came back from reflection is the name that's actually in
    the database, so it should not be elaborated further by a naming
    convention; :class:`sqlalchemy.schema.conv` is the marker that
    indicates this.

    """
    for const in table.constraints:
        if constraint_name_string(const.name) and not isinstance(
            const.name, conv
        ):
            const.name = conv(const.name)

    for idx in table.indexes:
        if constraint_name_string(idx.name) and not isinstance(idx.name, conv):
            idx.name = conv(idx.name)


def _name_type_bound_constraints(table: Table, source_table: Table) -> None:
    """Apply the naming convention on ``source_table.metadata`` to the
    type-bound constraints on ``table``.

    ``table`` here is a copy of ``source_table`` under a different name, so
    the convention is resolved against ``source_table`` in order for the
    ``%(table_name)s`` token to produce the real table name.

    """
    if not source_table.metadata.naming_convention:
        return

    for const in table.constraints:
        if not _is_type_bound(const):
            continue
        try:
            newname = _constraint_name_for_table(const, source_table)
        except sa_exc.InvalidRequestError:
            # the convention includes %(constraint_name)s but the type that
            # generates this constraint has no name of its own, so there's
            # nothing to interpolate.  the constraint is emitted without a
            # name, which is what it was before the naming convention was
            # taken into account here at all.
            colnames = ", ".join(
                c.name for c in _columns_for_constraint(const)
            )
            warn(
                f"The CHECK constraint generated by the type of column "
                f"'{source_table.name}.{colnames}' has no name, while the "
                f"naming convention in use requires one; the constraint is "
                f"emitted without a name.  Boolean and Enum datatypes should "
                f"be given a name, e.g. Boolean(name='...'), in order for "
                f"their CHECK constraints to participate fully in batch "
                f"mode."
            )
            continue
        if newname:
            const.name = newname


def _find_columns(clause):
    """locate Column objects within the given expression."""

    cols: set[ColumnElement[Any]] = set()
    traverse(clause, {}, {"column": cols.add})
    return cols


def _remove_column_from_collection(
    collection: ColumnCollection, column: Column[Any] | ColumnClause[Any]
) -> None:
    """remove a column from a ColumnCollection."""

    # workaround for older SQLAlchemy, remove the
    # same object that's present
    assert column.key is not None
    to_remove = collection[column.key]

    # SQLAlchemy 2.0 will use more ReadOnlyColumnCollection
    # (renamed from ImmutableColumnCollection)
    if hasattr(collection, "_immutable") or hasattr(collection, "_readonly"):
        collection._parent.remove(to_remove)
    else:
        collection.remove(to_remove)


def _textual_index_column(
    table: Table, text_: str | TextClause | ColumnElement[Any]
) -> ColumnElement[Any] | Column[Any]:
    """a workaround for the Index construct's severe lack of flexibility"""
    if isinstance(text_, str):
        c = Column(text_, sqltypes.NULLTYPE)
        table.append_column(c)
        return c
    elif isinstance(text_, TextClause):
        return _textual_index_element(table, text_)
    elif isinstance(text_, _textual_index_element):
        return _textual_index_column(table, text_.text)
    elif isinstance(text_, sql.ColumnElement):
        return _copy_expression(text_, table)
    else:
        raise ValueError("String or text() construct expected")


def _copy_expression(expression: _CE, target_table: Table) -> _CE:
    def replace(col):
        if (
            isinstance(col, Column)
            and col.table is not None
            and col.table is not target_table
        ):
            if col.name in target_table.c:
                return target_table.c[col.name]
            else:
                c = _copy(col)
                target_table.append_column(c)
                return c
        else:
            return None

    return visitors.replacement_traverse(  # type: ignore[call-overload]
        expression, {}, replace
    )


class _textual_index_element(sql.ColumnElement):
    """Wrap around a sqlalchemy text() construct in such a way that
    we appear like a column-oriented SQL expression to an Index
    construct.

    The issue here is that currently the Postgresql dialect, the biggest
    recipient of functional indexes, keys all the index expressions to
    the corresponding column expressions when rendering CREATE INDEX,
    so the Index we create here needs to have a .columns collection that
    is the same length as the .expressions collection.  Ultimately
    SQLAlchemy should support text() expressions in indexes.

    See SQLAlchemy issue 3174.

    """

    __visit_name__ = "_textual_idx_element"

    def __init__(self, table: Table, text: TextClause) -> None:
        self.table = table
        self.text = text
        self.key = text.text
        self.fake_column = schema.Column(self.text.text, sqltypes.NULLTYPE)
        table.append_column(self.fake_column)

    def get_children(self, **kw):
        return [self.fake_column]


@compiles(_textual_index_element)
def _render_textual_index_column(
    element: _textual_index_element, compiler: SQLCompiler, **kw
) -> str:
    return compiler.process(element.text, **kw)


class _literal_bindparam(BindParameter):
    pass


@compiles(_literal_bindparam)
def _render_literal_bindparam(
    element: _literal_bindparam, compiler: SQLCompiler, **kw
) -> str:
    return compiler.render_literal_bindparam(element, **kw)


def _get_constraint_final_name(
    constraint: Index | Constraint, dialect: Dialect | None
) -> str | None:
    if constraint.name is None:
        return None
    assert dialect is not None
    # make use of SQLAlchemy API to return what would be the final
    # compiled form of the name for this dialect.
    return dialect.identifier_preparer.format_constraint(
        constraint, _alembic_quote=False
    )


def _constraint_is_named(
    constraint: Constraint | Index, dialect: Dialect | None
) -> bool:
    if constraint.name is None:
        return False
    assert dialect is not None
    name = dialect.identifier_preparer.format_constraint(
        constraint, _alembic_quote=False
    )
    return name is not None


def is_expression_index(index: Index) -> bool:
    for expr in index.expressions:
        if is_expression(expr):
            return True
    return False


def is_expression(expr: Any) -> bool:
    while isinstance(expr, UnaryExpression):
        expr = expr.element
    if not isinstance(expr, ColumnClause) or expr.is_literal:
        return True
    return False


def _inherit_schema_deprecated() -> bool:
    # at some point in 2.1 inherit_schema was replaced with a property
    # so that's preset at the class level, while before it wasn't.
    return sqla_2_1 and hasattr(sqltypes.Enum, "inherit_schema")


def all_table_check_constraints(table: Table) -> set[CheckConstraint]:
    """Returns all check constraint including type bound ones and
    those on columns."""
    candidates = chain(
        table.constraints, *(c.constraints for c in table.columns)
    )
    return {ck for ck in candidates if isinstance(ck, CheckConstraint)}