"""This file includes annotation-sensitive tests while having
``from __future__ import annotations`` in effect.

Only tests that don't have an equivalent in ``test_typed_mapping`` are
specified here. All test from ``test_typed_mapping`` are copied over to
the ``test_tm_future_annotations_sync`` by the ``sync_test_file`` script.
"""

from __future__ import annotations

import enum
from typing import ClassVar
from typing import Dict
from typing import List
from typing import Optional
from typing import TYPE_CHECKING
from typing import TypeVar
import uuid

from sqlalchemy import exc
from sqlalchemy import ForeignKey
from sqlalchemy import inspect
from sqlalchemy import Integer
from sqlalchemy import select
from sqlalchemy import testing
from sqlalchemy import Uuid
import sqlalchemy.orm
from sqlalchemy.orm import attribute_keyed_dict
from sqlalchemy.orm import KeyFuncDict
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.orm import relationship
from sqlalchemy.orm.util import _cleanup_mapped_str_annotation
from sqlalchemy.sql import sqltypes
from sqlalchemy.testing import eq_
from sqlalchemy.testing import expect_raises_message
from sqlalchemy.testing import fixtures
from sqlalchemy.testing import is_
from sqlalchemy.testing import is_true
from .test_typed_mapping import expect_annotation_syntax_error
from .test_typed_mapping import MappedColumnTest as _MappedColumnTest
from .test_typed_mapping import RelationshipLHSTest as _RelationshipLHSTest


_R = TypeVar("_R")

M = Mapped


class M3:
    pass


class AnnoUtilTest(fixtures.TestBase):
    @testing.combinations(
        ("Mapped[Address]", 'Mapped["Address"]'),
        ('Mapped["Address"]', 'Mapped["Address"]'),
        ("Mapped['Address']", "Mapped['Address']"),
        ("Mapped[Address | None]", 'Mapped["Address | None"]'),
        ("Mapped[None | Address]", 'Mapped["None | Address"]'),
        ('Mapped["Address | None"]', 'Mapped["Address | None"]'),
        ("Mapped['None | Address']", "Mapped['None | Address']"),
        ('Mapped["Address" | "None"]', 'Mapped["Address" | "None"]'),
        ('Mapped["None" | "Address"]', 'Mapped["None" | "Address"]'),
        ("Mapped[A_]", 'Mapped["A_"]'),
        ("Mapped[_TypingLiteral]", 'Mapped["_TypingLiteral"]'),
        ("Mapped[datetime.datetime]", 'Mapped["datetime.datetime"]'),
        ("Mapped[List[Edge]]", 'Mapped[List["Edge"]]'),
        (
            "Mapped[collections.abc.MutableSequence[B]]",
            'Mapped[collections.abc.MutableSequence["B"]]',
        ),
        ("Mapped[typing.Sequence[B]]", 'Mapped[typing.Sequence["B"]]'),
        ("Mapped[dict[str, str]]", 'Mapped[dict["str", "str"]]'),
        ("Mapped[Dict[str, str]]", 'Mapped[Dict["str", "str"]]'),
        ("Mapped[list[str]]", 'Mapped[list["str"]]'),
        ("Mapped[dict[str, str] | None]", "Mapped[dict[str, str] | None]"),
        ("Mapped[Optional[anno_str_mc]]", 'Mapped[Optional["anno_str_mc"]]'),
        (
            "Mapped[Optional[Dict[str, str]]]",
            'Mapped[Optional[Dict["str", "str"]]]',
        ),
        (
            "Mapped[Optional[Union[Decimal, float]]]",
            'Mapped[Optional[Union["Decimal", "float"]]]',
        ),
        (
            "Mapped[Optional[Union[list[int], list[str]]]]",
            "Mapped[Optional[Union[list[int], list[str]]]]",
        ),
        ("Mapped[TestType[str]]", 'Mapped[TestType["str"]]'),
        ("Mapped[TestType[str, str]]", 'Mapped[TestType["str", "str"]]'),
        ("Mapped[Union[A, None]]", 'Mapped[Union["A", "None"]]'),
        ("Mapped[Union[Decimal, float]]", 'Mapped[Union["Decimal", "float"]]'),
        (
            "Mapped[Union[Decimal, float, None]]",
            'Mapped[Union["Decimal", "float", "None"]]',
        ),
        (
            "Mapped[Union[Dict[str, str], None]]",
            "Mapped[Union[Dict[str, str], None]]",
        ),
        ("Mapped[Union[float, Decimal]]", 'Mapped[Union["float", "Decimal"]]'),
        (
            "Mapped[Union[list[int], list[str]]]",
            "Mapped[Union[list[int], list[str]]]",
        ),
        (
            "Mapped[Union[list[int], list[str], None]]",
            "Mapped[Union[list[int], list[str], None]]",
        ),
        (
            "Mapped[Union[None, Dict[str, str]]]",
            "Mapped[Union[None, Dict[str, str]]]",
        ),
        (
            "Mapped[Union[None, list[int], list[str]]]",
            "Mapped[Union[None, list[int], list[str]]]",
        ),
        ("Mapped[A | None]", 'Mapped["A | None"]'),
        ("Mapped[Decimal | float]", 'Mapped["Decimal | float"]'),
        ("Mapped[Decimal | float | None]", 'Mapped["Decimal | float | None"]'),
        (
            "Mapped[list[int] | list[str] | None]",
            "Mapped[list[int] | list[str] | None]",
        ),
        ("Mapped[None | dict[str, str]]", "Mapped[None | dict[str, str]]"),
        (
            "Mapped[None | list[int] | list[str]]",
            "Mapped[None | list[int] | list[str]]",
        ),
    )
    def test_cleanup_mapped_str_annotation(self, given, expected):
        eq_(_cleanup_mapped_str_annotation(given, __name__), expected)


class MappedColumnTest(_MappedColumnTest):
    def test_fully_qualified_mapped_name(self, decl_base):
        """test #8853, regression caused by #8759 ;)


        See same test in test_abs_import_only

        """

        class Foo(decl_base):
            __tablename__ = "foo"

            id: sqlalchemy.orm.Mapped[int] = mapped_column(primary_key=True)

            data: sqlalchemy.orm.Mapped[int] = mapped_column()

            data2: sqlalchemy.orm.Mapped[int]

        self.assert_compile(
            select(Foo), "SELECT foo.id, foo.data, foo.data2 FROM foo"
        )

    def test_indirect_mapped_name_module_level(self, decl_base):
        """test #8759


        Note that M by definition has to be at the module level to be
        valid, and not locally declared here, this is in accordance with
        mypy::


            def make_class() -> None:
                ll = list

                x: ll[int] = [1, 2, 3]

        Will return::

            $ mypy test3.py
            test3.py:4: error: Variable "ll" is not valid as a type  [valid-type]
            test3.py:4: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
            Found 1 error in 1 file (checked 1 source file)

        Whereas the correct form is::

            ll = list


            def make_class() -> None:

                x: ll[int] = [1, 2, 3]

        """  # noqa: E501

        class Foo(decl_base):
            __tablename__ = "foo"

            id: M[int] = mapped_column(primary_key=True)

            data: M[int] = mapped_column()

            data2: M[int]

        self.assert_compile(
            select(Foo), "SELECT foo.id, foo.data, foo.data2 FROM foo"
        )

    def test_type_favors_outer(self, decl_base):
        """test #10899, that we maintain favoring outer names vs. inner.
        this is for backwards compatibility as well as what people
        usually expect regarding the names of attributes in the class.

        """

        class User(decl_base):
            __tablename__ = "user"

            id: Mapped[int] = mapped_column(primary_key=True)
            uuid: Mapped[uuid.UUID] = mapped_column()

        is_true(isinstance(User.__table__.c.uuid.type, sqltypes.Uuid))

    def test_type_inline_cls_qualified(self, decl_base):
        """test #10899, where we test that we can refer to the class name
        directly to refer to class-bound elements.

        """

        class User(decl_base):
            __tablename__ = "user"

            class Role(enum.Enum):
                admin = "admin"
                user = "user"

            id: Mapped[int] = mapped_column(primary_key=True)
            role: Mapped[User.Role]

        is_true(isinstance(User.__table__.c.role.type, sqltypes.Enum))
        eq_(User.__table__.c.role.type.length, 5)
        is_(User.__table__.c.role.type.enum_class, User.Role)

    def test_type_inline_disambiguate(self, decl_base):
        """test #10899, where we test that we can refer to an inner name
        that's not in conflict directly without qualification.

        """

        class User(decl_base):
            __tablename__ = "user"

            class Role(enum.Enum):
                admin = "admin"
                user = "user"

            id: Mapped[int] = mapped_column(primary_key=True)
            role: Mapped[Role]

        is_true(isinstance(User.__table__.c.role.type, sqltypes.Enum))
        eq_(User.__table__.c.role.type.length, 5)
        is_(User.__table__.c.role.type.enum_class, User.Role)
        eq_(User.__table__.c.role.type.name, "role")  # and not 'enum'

    def test_type_inner_can_be_qualified(self, decl_base):
        """test #10899, same test as that of Role, using it to qualify against
        a global variable with the same name.

        """

        global SomeGlobalName
        SomeGlobalName = None

        class User(decl_base):
            __tablename__ = "user"

            class SomeGlobalName(enum.Enum):
                admin = "admin"
                user = "user"

            id: Mapped[int] = mapped_column(primary_key=True)
            role: Mapped[User.SomeGlobalName]

        is_true(isinstance(User.__table__.c.role.type, sqltypes.Enum))
        eq_(User.__table__.c.role.type.length, 5)
        is_(User.__table__.c.role.type.enum_class, User.SomeGlobalName)

    def test_indirect_mapped_name_local_level(self, decl_base):
        """test #8759.

        this should raise an error.

        """

        M2 = Mapped

        with expect_raises_message(
            exc.ArgumentError,
            r"Could not interpret annotation M2\[int\].  Check that it "
            "uses names that are correctly imported at the module level.",
        ):

            class Foo(decl_base):
                __tablename__ = "foo"

                id: M2[int] = mapped_column(primary_key=True)

                data2: M2[int]

    def test_indirect_mapped_name_itswrong(self, decl_base):
        """test #8759.

        this should raise an error.

        """

        with expect_annotation_syntax_error("Foo.id"):

            class Foo(decl_base):
                __tablename__ = "foo"

                id: M3[int] = mapped_column(primary_key=True)

                data2: M3[int]

    def test_typ_not_in_cls_namespace(self, decl_base):
        """test #8742.

        This tests that when types are resolved, they use the ``__module__``
        of they class they are used within, not the mapped class.

        """

        class Mixin:
            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[uuid.UUID]

        class MyClass(Mixin, decl_base):
            # basically no type will be resolvable here
            __module__ = "some.module"
            __tablename__ = "mytable"

        is_(MyClass.id.expression.type._type_affinity, Integer)
        is_(MyClass.data.expression.type._type_affinity, Uuid)

    def test_dont_ignore_unresolvable(self, decl_base):
        """test #8888"""

        with expect_raises_message(
            exc.ArgumentError,
            r"Could not resolve all types within mapped annotation: "
            r"\"Mapped\[fake\]\".  Ensure all types are written correctly and "
            r"are imported within the module in use.",
        ):

            class A(decl_base):
                __tablename__ = "a"

                id: Mapped[int] = mapped_column(primary_key=True)
                data: Mapped[fake]  # noqa

    @testing.variation(
        "reference_type",
        [
            "plain",
            "plain_optional",
            "container_w_local_mapped",
            "container_w_remote_mapped",
        ],
    )
    def test_i_have_a_classvar_on_my_class(self, decl_base, reference_type):
        if reference_type.container_w_remote_mapped:

            class MyOtherClass(decl_base):
                __tablename__ = "myothertable"

                id: Mapped[int] = mapped_column(primary_key=True)

        class MyClass(decl_base):
            __tablename__ = "mytable"

            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[str] = mapped_column(default="some default")

            if reference_type.container_w_remote_mapped:
                status: ClassVar[Dict[str, MyOtherClass]]
            elif reference_type.container_w_local_mapped:
                status: ClassVar[Dict[str, MyClass]]
            elif reference_type.plain_optional:
                status: ClassVar[Optional[int]]
            elif reference_type.plain:
                status: ClassVar[int]

        m1 = MyClass(id=1, data=5)
        assert "status" not in inspect(m1).mapper.attrs


class MappedOneArg(KeyFuncDict[str, _R]):
    pass


class RelationshipLHSTest(_RelationshipLHSTest):
    def test_bidirectional_literal_annotations(self, decl_base):
        """test the 'string cleanup' function in orm/util.py, where
        we receive a string annotation like::

            "Mapped[List[B]]"

        Which then fails to evaluate because we don't have "B" yet.
        The annotation is converted on the fly to::

            'Mapped[List["B"]]'

        so that when we evaluated it, we get ``Mapped[List["B"]]`` and
        can extract "B" as a string.

        """

        class A(decl_base):
            __tablename__ = "a"

            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[str] = mapped_column()
            bs: Mapped[List[B]] = relationship(back_populates="a")

        class B(decl_base):
            __tablename__ = "b"
            id: Mapped[int] = mapped_column(Integer, primary_key=True)
            a_id: Mapped[int] = mapped_column(ForeignKey("a.id"))

            a: Mapped[A] = relationship(
                back_populates="bs", primaryjoin=a_id == A.id
            )

        a1 = A(data="data")
        b1 = B()
        a1.bs.append(b1)
        is_(a1, b1.a)

    def test_collection_class_dict_attr_mapped_collection_literal_annotations(
        self, decl_base
    ):
        class A(decl_base):
            __tablename__ = "a"

            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[str] = mapped_column()

            bs: Mapped[KeyFuncDict[str, B]] = relationship(  # noqa: F821
                collection_class=attribute_keyed_dict("name")
            )

        class B(decl_base):
            __tablename__ = "b"
            id: Mapped[int] = mapped_column(Integer, primary_key=True)
            a_id: Mapped[int] = mapped_column(ForeignKey("a.id"))
            name: Mapped[str] = mapped_column()

        self._assert_dict(A, B)

    def test_collection_cls_attr_mapped_collection_dbl_literal_annotations(
        self, decl_base
    ):
        class A(decl_base):
            __tablename__ = "a"

            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[str] = mapped_column()

            bs: Mapped[KeyFuncDict[str, "B"]] = relationship(  # noqa: F821
                collection_class=attribute_keyed_dict("name")
            )

        class B(decl_base):
            __tablename__ = "b"
            id: Mapped[int] = mapped_column(Integer, primary_key=True)
            a_id: Mapped[int] = mapped_column(ForeignKey("a.id"))
            name: Mapped[str] = mapped_column()

        self._assert_dict(A, B)

    def test_collection_cls_not_locatable(self, decl_base):
        class MyCollection(KeyFuncDict):
            pass

        with expect_raises_message(
            exc.ArgumentError,
            r"Could not interpret annotation Mapped\[MyCollection\['B'\]\].",
        ):

            class A(decl_base):
                __tablename__ = "a"

                id: Mapped[int] = mapped_column(primary_key=True)
                data: Mapped[str] = mapped_column()

                bs: Mapped[MyCollection["B"]] = relationship(  # noqa: F821
                    collection_class=attribute_keyed_dict("name")
                )

    def test_collection_cls_one_arg(self, decl_base):
        class A(decl_base):
            __tablename__ = "a"

            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[str] = mapped_column()

            bs: Mapped[MappedOneArg["B"]] = relationship(  # noqa: F821
                collection_class=attribute_keyed_dict("name")
            )

        class B(decl_base):
            __tablename__ = "b"
            id: Mapped[int] = mapped_column(Integer, primary_key=True)
            a_id: Mapped[int] = mapped_column(ForeignKey("a.id"))
            name: Mapped[str] = mapped_column()

        self._assert_dict(A, B)

    def _assert_dict(self, A, B):
        A.registry.configure()

        a1 = A()
        b1 = B(name="foo")

        # collection appender on MappedCollection
        a1.bs.set(b1)

        is_(a1.bs["foo"], b1)

    def test_indirect_name_relationship_arg_override(self, decl_base):
        """test #8759

        in this test we assume a case where the type for the Mapped annnotation
        a. has to be a different name than the actual class name and
        b. cannot be imported outside of TYPE CHECKING.  user will then put
        the real name inside of relationship().  we have to succeed even though
        we can't resolve the annotation.

        """

        class B(decl_base):
            __tablename__ = "b"
            id: Mapped[int] = mapped_column(Integer, primary_key=True)
            a_id: Mapped[int] = mapped_column(ForeignKey("a.id"))

        if TYPE_CHECKING:
            BNonExistent = B

        class A(decl_base):
            __tablename__ = "a"

            id: Mapped[int] = mapped_column(primary_key=True)
            data: Mapped[str] = mapped_column()

            bs: Mapped[List[BNonExistent]] = relationship("B")

        self.assert_compile(
            select(A).join(A.bs),
            "SELECT a.id, a.data FROM a JOIN b ON a.id = b.a_id",
        )