from sqlalchemy import Integer, String, ForeignKey, func, desc, and_, or_
from sqlalchemy.orm import interfaces, relationship, mapper, \
    clear_mappers, create_session, joinedload, joinedload_all, \
    subqueryload, subqueryload_all, polymorphic_union, aliased,\
    class_mapper, with_polymorphic
from sqlalchemy import exc as sa_exc
from sqlalchemy.engine import default

from sqlalchemy.testing import AssertsCompiledSQL, fixtures
from sqlalchemy import testing
from sqlalchemy.testing.schema import Table, Column
from sqlalchemy.testing import assert_raises, eq_

from ._poly_fixtures import Company, Person, Engineer, Manager, Boss, \
    Machine, Paperwork, _PolymorphicFixtureBase, _Polymorphic,\
    _PolymorphicPolymorphic, _PolymorphicUnions, _PolymorphicJoins,\
    _PolymorphicAliasedJoins


class _WithPolymorphicBase(_PolymorphicFixtureBase):
    def test_join_base_to_sub(self):
        sess = create_session()
        pa = with_polymorphic(Person, [Engineer])

        def go():
            eq_(sess.query(pa)
                    .filter(pa.Engineer.primary_language == 'java').all(),
                self._emps_wo_relationships_fixture()[0:1])
        self.assert_sql_count(testing.db, go, 1)

    def test_col_expression_base_plus_two_subs(self):
        sess = create_session()
        pa = with_polymorphic(Person, [Engineer, Manager])

        eq_(sess.query(
            pa.name, pa.Engineer.primary_language,
            pa.Manager.manager_name).filter(
            or_(
                pa.Engineer.primary_language == 'java', pa.Manager.
                manager_name
                == 'dogbert')).order_by(pa.Engineer.type).all(),
            [('dilbert', 'java', None),
             ('dogbert', None, 'dogbert'), ])

    def test_join_to_join_entities(self):
        sess = create_session()
        pa = with_polymorphic(Person, [Engineer])
        pa_alias = with_polymorphic(Person, [Engineer], aliased=True)

        eq_(
            [(p1.name, type(p1), p2.name, type(p2)) for (p1, p2) in sess.query(
                pa, pa_alias
            ).join(pa_alias,
                   or_(
                       pa.Engineer.primary_language ==
                       pa_alias.Engineer.primary_language,
                       and_(
                           pa.Engineer.primary_language == None,  # noqa
                           pa_alias.Engineer.primary_language == None,
                           pa.person_id > pa_alias.person_id
                       ))
                   ).order_by(pa.name, pa_alias.name)],
            [
                ('dilbert', Engineer, 'dilbert', Engineer),
                ('dogbert', Manager, 'pointy haired boss', Boss),
                ('vlad', Engineer, 'vlad', Engineer),
                ('wally', Engineer, 'wally', Engineer)
            ]
        )

    def test_join_to_join_columns(self):
        sess = create_session()
        pa = with_polymorphic(Person, [Engineer])
        pa_alias = with_polymorphic(Person, [Engineer], aliased=True)

        eq_(
            [row for row in sess.query(
                pa.name, pa.Engineer.primary_language,
                pa_alias.name, pa_alias.Engineer.primary_language
            ).join(pa_alias,
                   or_(
                       pa.Engineer.primary_language ==
                       pa_alias.Engineer.primary_language,
                       and_(
                           pa.Engineer.primary_language == None,  # noqa
                           pa_alias.Engineer.primary_language == None,
                           pa.person_id > pa_alias.person_id
                       ))
                   ).order_by(pa.name, pa_alias.name)],
            [
                ('dilbert', 'java', 'dilbert', 'java'),
                ('dogbert', None, 'pointy haired boss', None),
                ('vlad', 'cobol', 'vlad', 'cobol'),
                ('wally', 'c++', 'wally', 'c++')
            ]
        )


class PolymorphicTest(_WithPolymorphicBase, _Polymorphic):
    pass


class PolymorphicPolymorphicTest(_WithPolymorphicBase,
                                 _PolymorphicPolymorphic):
    pass


class PolymorphicUnionsTest(_WithPolymorphicBase, _PolymorphicUnions):
    pass


class PolymorphicAliasedJoinsTest(_WithPolymorphicBase,
                                  _PolymorphicAliasedJoins):
    pass


class PolymorphicJoinsTest(_WithPolymorphicBase, _PolymorphicJoins):
    pass
