import sqlalchemy as sa
from sqlalchemy import testing
from sqlalchemy import util
from sqlalchemy.orm import create_session
from sqlalchemy.orm import mapper
from sqlalchemy.orm import relationship
from sqlalchemy.testing import assert_raises_message
from sqlalchemy.testing import eq_
from test.orm import _fixtures


class DefaultStrategyOptionsTest(_fixtures.FixtureTest):
    def _assert_fully_loaded(self, users):
        # verify everything loaded, with no additional sql needed
        def go():
            # comparison with no additional sql
            eq_(users, self.static.user_all_result)
            # keywords are not part of self.static.user_all_result, so
            # verify all the item keywords were loaded, with no more sql.
            # 'any' verifies at least some items have keywords; we build
            # a list for any([...]) instead of any(...) to prove we've
            # iterated all the items with no sql.
            f = util.flatten_iterator
            assert any(
                [
                    i.keywords
                    for i in f([o.items for o in f([u.orders for u in users])])
                ]
            )

        self.assert_sql_count(testing.db, go, 0)

    def _assert_addresses_loaded(self, users):
        # verify all the addresses were joined loaded with no more sql
        def go():
            for u, static in zip(users, self.static.user_all_result):
                eq_(u.addresses, static.addresses)

        self.assert_sql_count(testing.db, go, 0)

    def _downgrade_fixture(self):
        (
            users,
            Keyword,
            items,
            order_items,
            orders,
            Item,
            User,
            Address,
            keywords,
            item_keywords,
            Order,
            addresses,
        ) = (
            self.tables.users,
            self.classes.Keyword,
            self.tables.items,
            self.tables.order_items,
            self.tables.orders,
            self.classes.Item,
            self.classes.User,
            self.classes.Address,
            self.tables.keywords,
            self.tables.item_keywords,
            self.classes.Order,
            self.tables.addresses,
        )

        mapper(Address, addresses)

        mapper(Keyword, keywords)

        mapper(
            Item,
            items,
            properties=dict(
                keywords=relationship(
                    Keyword,
                    secondary=item_keywords,
                    lazy="subquery",
                    order_by=item_keywords.c.keyword_id,
                )
            ),
        )

        mapper(
            Order,
            orders,
            properties=dict(
                items=relationship(
                    Item,
                    secondary=order_items,
                    lazy="subquery",
                    order_by=order_items.c.item_id,
                )
            ),
        )

        mapper(
            User,
            users,
            properties=dict(
                addresses=relationship(
                    Address, lazy="joined", order_by=addresses.c.id
                ),
                orders=relationship(
                    Order, lazy="joined", order_by=orders.c.id
                ),
            ),
        )

        return create_session()

    def _upgrade_fixture(self):
        (
            users,
            Keyword,
            items,
            order_items,
            orders,
            Item,
            User,
            Address,
            keywords,
            item_keywords,
            Order,
            addresses,
        ) = (
            self.tables.users,
            self.classes.Keyword,
            self.tables.items,
            self.tables.order_items,
            self.tables.orders,
            self.classes.Item,
            self.classes.User,
            self.classes.Address,
            self.tables.keywords,
            self.tables.item_keywords,
            self.classes.Order,
            self.tables.addresses,
        )

        mapper(Address, addresses)

        mapper(Keyword, keywords)

        mapper(
            Item,
            items,
            properties=dict(
                keywords=relationship(
                    Keyword,
                    secondary=item_keywords,
                    lazy="select",
                    order_by=item_keywords.c.keyword_id,
                )
            ),
        )

        mapper(
            Order,
            orders,
            properties=dict(
                items=relationship(
                    Item,
                    secondary=order_items,
                    lazy=True,
                    order_by=order_items.c.item_id,
                )
            ),
        )

        mapper(
            User,
            users,
            properties=dict(
                addresses=relationship(
                    Address, lazy=True, order_by=addresses.c.id
                ),
                orders=relationship(Order, order_by=orders.c.id),
            ),
        )

        return create_session()

    def test_downgrade_baseline(self):
        """Mapper strategy defaults load as expected
        (compare to rest of DefaultStrategyOptionsTest downgrade tests)."""
        sess = self._downgrade_fixture()
        users = []

        # test _downgrade_fixture mapper defaults, 3 queries (2 subquery
        # loads).
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 3)

        # all loaded with no additional sql
        self._assert_fully_loaded(users)

    def test_disable_eagerloads(self):
        """Mapper eager load strategy defaults can be shut off
        with enable_eagerloads(False)."""

        # While this isn't testing a mapper option, it is included
        # as baseline reference for how XYZload('*') option
        # should work, namely, it shouldn't affect later queries
        # (see other test_select_s)
        sess = self._downgrade_fixture()
        users = []

        # demonstrate that enable_eagerloads loads with only 1 sql
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .enable_eagerloads(False)
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)

        # demonstrate that users[0].orders must now be loaded with 3 sql
        # (need to lazyload, and 2 subquery: 3 total)
        def go():
            users[0].orders

        self.assert_sql_count(testing.db, go, 3)

    def test_last_one_wins(self):
        sess = self._downgrade_fixture()
        users = []

        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.subqueryload("*"))
                .options(sa.orm.joinedload(self.classes.User.addresses))
                .options(sa.orm.lazyload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)

        # verify all the addresses were joined loaded (no more sql)
        self._assert_addresses_loaded(users)

    def test_star_must_be_alone(self):
        sess = self._downgrade_fixture()
        User = self.classes.User
        opt = sa.orm.subqueryload("*", User.addresses)
        assert_raises_message(
            sa.exc.ArgumentError,
            "Wildcard token cannot be followed by another entity",
            sess.query(User).options(opt)._compile_context,
        )

    def test_global_star_ignored_no_entities_unbound(self):
        sess = self._downgrade_fixture()
        User = self.classes.User
        opt = sa.orm.lazyload("*")
        q = sess.query(User.name).options(opt)
        eq_(q.all(), [("jack",), ("ed",), ("fred",), ("chuck",)])

    def test_global_star_ignored_no_entities_bound(self):
        sess = self._downgrade_fixture()
        User = self.classes.User
        opt = sa.orm.Load(User).lazyload("*")
        q = sess.query(User.name).options(opt)
        eq_(q.all(), [("jack",), ("ed",), ("fred",), ("chuck",)])

    def test_select_with_joinedload(self):
        """Mapper load strategy defaults can be downgraded with
        lazyload('*') option, while explicit joinedload() option
        is still honored"""
        sess = self._downgrade_fixture()
        users = []

        # lazyload('*') shuts off 'orders' subquery: only 1 sql
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.lazyload("*"))
                .options(sa.orm.joinedload(self.classes.User.addresses))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)

        # verify all the addresses were joined loaded (no more sql)
        self._assert_addresses_loaded(users)

        # users[0] has orders, which need to lazy load, and 2 subquery:
        # (same as with test_disable_eagerloads): 3 total sql
        def go():
            users[0].orders

        self.assert_sql_count(testing.db, go, 3)

    def test_select_with_subqueryload(self):
        """Mapper load strategy defaults can be downgraded with
        lazyload('*') option, while explicit subqueryload() option
        is still honored"""
        sess = self._downgrade_fixture()
        users = []

        # now test 'default_strategy' option combined with 'subquery'
        # shuts off 'addresses' load AND orders.items load: 2 sql expected
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.lazyload("*"))
                .options(sa.orm.subqueryload(self.classes.User.orders))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 2)

        # Verify orders have already been loaded: 0 sql
        def go():
            for u, static in zip(users, self.static.user_all_result):
                assert len(u.orders) == len(static.orders)

        self.assert_sql_count(testing.db, go, 0)

        # Verify lazyload('*') prevented orders.items load
        # users[0].orders[0] has 3 items, each with keywords: 2 sql
        # ('items' and 'items.keywords' subquery)
        def go():
            for i in users[0].orders[0].items:
                i.keywords

        self.assert_sql_count(testing.db, go, 2)

        # lastly, make sure they actually loaded properly
        eq_(users, self.static.user_all_result)

    def test_noload_with_joinedload(self):
        """Mapper load strategy defaults can be downgraded with
        noload('*') option, while explicit joinedload() option
        is still honored"""
        sess = self._downgrade_fixture()
        users = []

        # test noload('*') shuts off 'orders' subquery, only 1 sql
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.noload("*"))
                .options(sa.orm.joinedload(self.classes.User.addresses))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)

        # verify all the addresses were joined loaded (no more sql)
        self._assert_addresses_loaded(users)

        # User.orders should have loaded "noload" (meaning [])
        def go():
            for u in users:
                assert u.orders == []

        self.assert_sql_count(testing.db, go, 0)

    def test_noload_with_subqueryload(self):
        """Mapper load strategy defaults can be downgraded with
        noload('*') option, while explicit subqueryload() option
        is still honored"""
        sess = self._downgrade_fixture()
        users = []

        # test noload('*') option combined with subqueryload()
        # shuts off 'addresses' load AND orders.items load: 2 sql expected
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.noload("*"))
                .options(sa.orm.subqueryload(self.classes.User.orders))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 2)

        def go():
            # Verify orders have already been loaded: 0 sql
            for u, static in zip(users, self.static.user_all_result):
                assert len(u.orders) == len(static.orders)
            # Verify noload('*') prevented orders.items load
            # and set 'items' to []
            for u in users:
                for o in u.orders:
                    assert o.items == []

        self.assert_sql_count(testing.db, go, 0)

    def test_joined(self):
        """Mapper load strategy defaults can be upgraded with
        joinedload('*') option."""
        sess = self._upgrade_fixture()
        users = []

        # test upgrade all to joined: 1 sql
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.joinedload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)

        # verify everything loaded, with no additional sql needed
        self._assert_fully_loaded(users)

    def test_joined_path_wildcards(self):
        sess = self._upgrade_fixture()
        users = []

        # test upgrade all to joined: 1 sql
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.joinedload(".*"))
                .options(sa.orm.joinedload("addresses.*"))
                .options(sa.orm.joinedload("orders.*"))
                .options(sa.orm.joinedload("orders.items.*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)
        self._assert_fully_loaded(users)

    def test_joined_with_lazyload(self):
        """Mapper load strategy defaults can be upgraded with
        joinedload('*') option, while explicit lazyload() option
        is still honored"""
        sess = self._upgrade_fixture()
        users = []

        # test joined all but 'keywords': upgraded to 1 sql
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.lazyload("orders.items.keywords"))
                .options(sa.orm.joinedload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 1)

        # everything (but keywords) loaded ok
        # (note self.static.user_all_result contains no keywords)
        def go():
            eq_(users, self.static.user_all_result)

        self.assert_sql_count(testing.db, go, 0)

        # verify the items were loaded, while item.keywords were not
        def go():
            # redundant with last test, but illustrative
            users[0].orders[0].items[0]

        self.assert_sql_count(testing.db, go, 0)

        def go():
            users[0].orders[0].items[0].keywords

        self.assert_sql_count(testing.db, go, 1)

    def test_joined_with_subqueryload(self):
        """Mapper load strategy defaults can be upgraded with
        joinedload('*') option, while explicit subqueryload() option
        is still honored"""
        sess = self._upgrade_fixture()
        users = []

        # test upgrade all but 'addresses', which is subquery loaded (2 sql)
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.subqueryload(self.classes.User.addresses))
                .options(sa.orm.joinedload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 2)

        # verify everything loaded, with no additional sql needed
        self._assert_fully_loaded(users)

    def test_subquery(self):
        """Mapper load strategy defaults can be upgraded with
        subqueryload('*') option."""
        sess = self._upgrade_fixture()
        users = []

        # test upgrade all to subquery: 1 sql + 4 relationships = 5
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.subqueryload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 5)

        # verify everything loaded, with no additional sql needed
        self._assert_fully_loaded(users)

    def test_subquery_path_wildcards(self):
        sess = self._upgrade_fixture()
        users = []

        # test upgrade all to subquery: 1 sql + 4 relationships = 5
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.subqueryload(".*"))
                .options(sa.orm.subqueryload("addresses.*"))
                .options(sa.orm.subqueryload("orders.*"))
                .options(sa.orm.subqueryload("orders.items.*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 5)

        # verify everything loaded, with no additional sql needed
        self._assert_fully_loaded(users)

    def test_subquery_with_lazyload(self):
        """Mapper load strategy defaults can be upgraded with
        subqueryload('*') option, while explicit lazyload() option
        is still honored"""
        sess = self._upgrade_fixture()
        users = []

        # test subquery all but 'keywords' (1 sql + 3 relationships = 4)
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.lazyload("orders.items.keywords"))
                .options(sa.orm.subqueryload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 4)

        # no more sql
        # (note self.static.user_all_result contains no keywords)
        def go():
            eq_(users, self.static.user_all_result)

        self.assert_sql_count(testing.db, go, 0)

        # verify the item.keywords were not loaded
        def go():
            users[0].orders[0].items[0]

        self.assert_sql_count(testing.db, go, 0)

        def go():
            users[0].orders[0].items[0].keywords

        self.assert_sql_count(testing.db, go, 1)

    def test_subquery_with_joinedload(self):
        """Mapper load strategy defaults can be upgraded with
        subqueryload('*') option, while multiple explicit
        joinedload() options are still honored"""
        sess = self._upgrade_fixture()
        users = []

        # test upgrade all but 'addresses' & 'orders', which are joinedloaded
        # (1 sql + items + keywords = 3)
        def go():
            users[:] = (
                sess.query(self.classes.User)
                .options(sa.orm.joinedload(self.classes.User.addresses))
                .options(sa.orm.joinedload(self.classes.User.orders))
                .options(sa.orm.subqueryload("*"))
                .order_by(self.classes.User.id)
                .all()
            )

        self.assert_sql_count(testing.db, go, 3)

        # verify everything loaded, with no additional sql needed
        self._assert_fully_loaded(users)