"""Test hooks & SQLAlchemy API integrations."""
import pytest

from sqlalchemy.orm import aliased, joinedload

from sqlalchemy_oso.session import (
    authorized_sessionmaker,
    scoped_session,
    AuthorizedSession,
)
from sqlalchemy_oso.compat import USING_SQLAlchemy_v1_3

from .models import User, Post
from .conftest import print_query


def log_queries():
    import logging

    logging.basicConfig()
    logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)


def test_authorize_query_no_access(engine, oso, fixture_data):
    session = AuthorizedSession(oso, "user", {Post: "action"}, bind=engine)
    query = session.query(Post)

    assert query.count() == 0


@pytest.mark.parametrize(
    "query",
    [
        lambda session: session.query(Post),
        lambda session: session.query(Post.contents, Post.id),
    ],
)
def test_authorize_query_basic(engine, oso, fixture_data, query):
    # TODO: copied from test_authorize_model_basic
    oso.load_str(
        """allow("user", "read", post: Post) if post.access_level = "public";
           allow("user", "write", post: Post) if post.access_level = "private";
           allow("admin", "read", _post: Post);
           allow("moderator", "read", post: Post) if
             (post.access_level = "private" or post.access_level = "public") and
             post.needs_moderation = true;"""
    )

    session = AuthorizedSession(oso, "user", {Post: "read"}, bind=engine)
    authorized = query(session)

    assert authorized.count() == 5
    assert authorized.all()[0].contents == "foo public post"
    assert authorized.all()[0].id == 0

    session = AuthorizedSession(oso, "user", {Post: "write"}, bind=engine)
    posts = query(session)

    assert posts.count() == 4
    assert posts.all()[0].contents == "foo private post"
    assert posts.all()[1].contents == "foo private post 2"

    session = AuthorizedSession(oso, "admin", {Post: "read"}, bind=engine)
    posts = query(session)
    assert posts.count() == 9

    session = AuthorizedSession(oso, "moderator", {Post: "read"}, bind=engine)
    posts = query(session)
    print_query(posts)
    assert posts.all()[0].contents == "private for moderation"
    assert posts.all()[1].contents == "public for moderation"

    session = AuthorizedSession(oso, "guest", {Post: "read"}, bind=engine)
    posts = query(session)
    assert posts.count() == 0


def test_authorize_query_multiple_types(engine, oso, fixture_data):
    """Test a query involving multiple models."""
    oso.load_str(
        """allow("user", "read", post: Post) if post.id = 1;
           allow("user", "read", user: User) if user.id = 0;
           allow("user", "read", user: User) if user.id = 1;
           allow("all_posts", "read", _: Post);"""
    )

    # Query two models. Only return authorized objects from each (no join).
    session = AuthorizedSession(oso, "user", {Post: "read", User: "read"}, bind=engine)
    authorized = session.query(Post, User)
    print_query(authorized)
    assert authorized.count() == 2
    assert authorized[0][0].id == 1
    assert authorized[0][1].id == 0
    assert authorized[1][1].id == 1

    # Query two models, with a join condition. Only return authorized objects that meet the join
    # condition.
    authorized = session.query(Post, User.username).join(User)
    print_query(authorized)
    assert authorized.count() == 1
    assert authorized[0][0].id == 1
    assert authorized[0][1] == "foo"

    # Join, but only return fields from one model.
    authorized = session.query(Post).join(User)

    # This one is odd... we don't return any fields from the User model,
    # so no authorization is applied for users.
    print_query(authorized)
    assert authorized.count() == 1

    # Another odd one.  We are joining on user, and filtering on fields on user.
    # But, no authorization filter is applied for user because no fields of user
    # are returned.
    # Could this leak data somehow? Maybe if users are allowed to filter arbitrary
    # values and see a count, but not retrieve the objects?
    session = AuthorizedSession(oso, "all_posts", {Post: "read"}, bind=engine)
    authorized = session.query(Post).join(User).filter(User.username == "admin_user")
    print_query(authorized)
    assert authorized.count() == 2

    # TODO (dhatch): What happens for aggregations?


@pytest.mark.xfail(USING_SQLAlchemy_v1_3, reason="Not supported by 1.3 events API.")
def test_authorize_query_joined_load(engine, oso, fixture_data):
    """Test a query involving multiple models."""
    oso.load_str(
        """allow("user", "read", post: Post) if post.id = 1;
           allow("user", "read", user: User) if user.id = 0;
           allow("user", "read", user: User) if user.id = 1;
           allow("all_posts", "read", _: Post);"""
    )

    session = AuthorizedSession(oso, "user", {Post: "read", User: "read"}, bind=engine)
    authorized = session.query(User).options(joinedload(User.posts))
    print_query(authorized)
    print(authorized[0].posts)
    assert len(authorized[0].posts) == 1


def test_alias(engine, oso, fixture_data):
    oso.load_str('allow("user", "read", post: Post) if post.id = 1;')
    session = AuthorizedSession(
        oso, user="user", checked_permissions={Post: "read"}, bind=engine
    )

    post_alias = aliased(Post)

    query = session.query(post_alias)

    # Fine with hooks if you don't authorize it.
    assert query.count() == 1


def test_authorized_sessionmaker_relationship(engine, oso, fixture_data):
    oso.load_str(
        """allow("user", "read", post: Post) if post.id = 1;
           # Post with creator id = 1
           allow("user", "read", post: Post) if post.id = 7;
           allow("user", "read", user: User) if user.id = 0;"""
    )

    Session = authorized_sessionmaker(
        get_oso=lambda: oso,
        get_user=lambda: "user",
        get_checked_permissions=lambda: {Post: "read", User: "read"},
        bind=engine,
    )

    session = Session()

    posts = session.query(Post)
    assert posts.count() == 2

    users = session.query(User)
    assert users.count() == 1

    post_1 = posts.get(1)
    # Authorized created by field.
    assert post_1.created_by == users.get(0)

    post_7 = posts.get(7)
    # created_by isn't actually none, but we can't see it
    assert post_7.created_by is None


def test_authorized_session_relationship(engine, oso, fixture_data):
    oso.load_str(
        """allow("user", "read", post: Post) if post.id = 1;
           # Post with creator id = 1
           allow("user", "read", post: Post) if post.id = 7;
           allow("user", "read", user: User) if user.id = 0;"""
    )

    session = AuthorizedSession(
        oso=oso,
        user="user",
        checked_permissions={Post: "read", User: "read"},
        bind=engine,
    )

    posts = session.query(Post)
    assert posts.count() == 2

    users = session.query(User)
    assert users.count() == 1

    post_1 = posts.get(1)
    # Authorized created by field.
    assert post_1.created_by == users.get(0)

    post_7 = posts.get(7)
    # created_by isn't actually none, but we can't see it
    assert post_7.created_by is None


def test_scoped_session_relationship(engine, oso, fixture_data):
    oso.load_str(
        """allow("user", "read", post: Post) if post.id = 1;
           # Post with creator id = 1
           allow("user", "read", post: Post) if post.id = 7;
           allow("user", "read", user: User) if user.id = 0;
           allow("other", "read", post: Post) if post.id = 3;
           allow("other", "write", post: Post) if post.id = 4;"""
    )

    data = {"user": "user", "checked_permissions": {Post: "read", User: "read"}}
    session = scoped_session(
        lambda: oso, lambda: data["user"], lambda: data["checked_permissions"]
    )
    session.configure(bind=engine)

    posts = session.query(Post)
    assert posts.count() == 2

    users = session.query(User)
    assert users.count() == 1

    post_1 = posts.get(1)
    # Authorized created by field.
    assert post_1.created_by == users.get(0)

    post_7 = posts.get(7)
    # created_by isn't actually none, but we can't see it
    assert post_7.created_by is None
    assert len(session.identity_map.values()) == 3

    data["user"] = "other"

    # Ensure this changed the session.
    assert len(session.identity_map.values()) == 0
    posts = session.query(Post)
    assert posts.count() == 1
    posts = posts.all()
    assert posts[0].id == 3
    assert len(session.identity_map.values()) == 1

    data["checked_permissions"] = {Post: "write", User: "write"}
    assert len(session.identity_map.values()) == 0
    posts = session.query(Post)
    assert posts.count() == 1
    posts = posts.all()
    assert posts[0].id == 4
    assert len(session.identity_map.values()) == 1

    # Change back to original.
    data = {"user": "user", "checked_permissions": {Post: "read", User: "read"}}
    assert len(session.identity_map.values()) == 3


@pytest.mark.xfail(reason="Implemented incorrectly initially. Fix")
def test_authorized_sessionmaker_user_change(engine, oso, fixture_data):
    """Ensure that query fails if the user changes."""
    oso.load_str('allow("user", "read", post: Post) if post.id = 1;')
    user = ["user"]

    Session = authorized_sessionmaker(
        get_oso=lambda: oso,
        get_user=lambda: user[0],
        get_checked_permissions=lambda: {Post: "read"},
        bind=engine,
    )

    session = Session()

    posts = session.query(Post).count()
    assert posts == 1

    user[0] = "moderator"

    with pytest.raises(Exception, match="user"):
        posts = session.query(Post).count()


def test_null_with_partial(engine, oso):
    oso.load_str("allow(_, _, post: Post) if post.contents = nil;")
    Session = authorized_sessionmaker(
        get_oso=lambda: oso,
        get_user=lambda: "user",
        get_checked_permissions=lambda: {Post: "read"},
        bind=engine,
    )
    posts = Session().query(Post)

    if USING_SQLAlchemy_v1_3:
        where_clause = " \nWHERE posts.contents IS NULL"
    else:
        # NOTE(gj): In constrast to the `Query.before_compile` event we listen
        # for in 1.3, the `Session.do_orm_execute` event we listen for in
        # SQLAlchemy 1.4 (unsurprisingly) happens at ORM execution time. Because
        # of this, the WHERE clauses we add as authorization constraints are
        # not applied when we (compile and) print out the query in the below
        # assertion but *are* applied when we actually interact with the ORM
        # when executing the `count()` method.
        where_clause = ""

    assert str(posts) == (
        "SELECT posts.id AS posts_id, posts.contents AS posts_contents, posts.title AS posts_title, "
        + "posts.access_level AS posts_access_level, posts.created_by_id AS posts_created_by_id, "
        + f"posts.needs_moderation AS posts_needs_moderation \nFROM posts{where_clause}"
    )
    assert posts.count() == 0


def test_regular_session(engine, oso, fixture_data):
    """Test that a regular session doesn't apply authorization."""
    from sqlalchemy.orm import Session

    session = Session(bind=engine)
    posts = session.query(Post)

    # No posts would be returned if authorization was applied.
    assert posts.count() == 9


def test_unconditional_policy_has_no_filter(engine, oso, fixture_data):
    oso.load_str('allow("user", "read", post: Post) if post.id = 1; allow(_, _, _);')
    session = AuthorizedSession(
        oso, user="user", checked_permissions={Post: "read"}, bind=engine
    )

    query = session.query(Post)

    if USING_SQLAlchemy_v1_3:
        where_clause = " \nWHERE 1 = 1"
    else:
        # NOTE(gj): In constrast to the `Query.before_compile` event we listen
        # for in 1.3, the `Session.do_orm_execute` event we listen for in
        # SQLAlchemy 1.4 (unsurprisingly) happens at ORM execution time. Because
        # of this, the WHERE clauses we add as authorization constraints are
        # not applied when we (compile and) print out the query in the below
        # assertion but *are* applied when we actually interact with the ORM
        # when executing the `count()` method.
        where_clause = ""

    assert str(query) == (
        "SELECT posts.id AS posts_id, posts.contents AS posts_contents, posts.title AS posts_title, "
        + "posts.access_level AS posts_access_level, posts.created_by_id AS posts_created_by_id, "
        + f"posts.needs_moderation AS posts_needs_moderation \nFROM posts{where_clause}"
    )
    assert query.count() == 9


def test_bakery_caching_for_AuthorizedSession(engine, oso, fixture_data):
    """Test that baked relationship queries don't lead to authorization bypasses
    for AuthorizedSession."""
    from sqlalchemy.orm import Session

    basic_session = Session(bind=engine)
    all_posts = basic_session.query(Post)
    assert all_posts.count() == 9
    first_post = all_posts[0]
    # Add related model query to the bakery cache.
    assert first_post.created_by.id == 0

    oso.load_str('allow("user", "read", post: Post) if post.id = 0;')

    # Baked queries disabled for sqlalchemy_oso.session.AuthorizedSession.
    authorized_session = AuthorizedSession(
        oso, user="user", checked_permissions={Post: "read"}, bind=engine
    )

    assert authorized_session.query(User).count() == 0

    authorized_posts = authorized_session.query(Post)
    assert authorized_posts.count() == 1
    first_authorized_post = authorized_posts[0]
    assert first_post.id == first_authorized_post.id

    # Should not be able to view the post's creator because there's no rule
    # permitting access to "read" users.
    assert first_authorized_post.created_by is None


def test_bakery_caching_for_authorized_sessionmaker(engine, oso, fixture_data):
    """Test that baked relationship queries don't lead to authorization bypasses
    for authorized_sessionmaker."""
    from sqlalchemy.orm import Session

    basic_session = Session(bind=engine)
    all_posts = basic_session.query(Post)
    assert all_posts.count() == 9
    first_post = all_posts[0]
    # Add related model query to the bakery cache.
    assert first_post.created_by.id == 0

    oso.load_str('allow("user", "read", post: Post) if post.id = 0;')

    # Baked queries disabled for sqlalchemy_oso.session.authorized_sessionmaker.
    authorized_session = authorized_sessionmaker(
        get_oso=lambda: oso,
        get_user=lambda: "user",
        get_checked_permissions=lambda: {Post: "read"},
        bind=engine,
    )()

    assert authorized_session.query(User).count() == 0

    authorized_posts = authorized_session.query(Post)
    assert authorized_posts.count() == 1
    first_authorized_post = authorized_posts[0]
    assert first_post.id == first_authorized_post.id

    # Should not be able to view the post's creator because there's no rule
    # permitting access to "read" users.
    assert first_authorized_post.created_by is None


def test_bakery_caching_for_scoped_session(engine, oso, fixture_data):
    """Test that baked relationship queries don't lead to authorization bypasses
    for scoped_session."""
    from sqlalchemy.orm import Session

    basic_session = Session(bind=engine)
    all_posts = basic_session.query(Post)
    assert all_posts.count() == 9
    first_post = all_posts[0]
    # Add related model query to the bakery cache.
    assert first_post.created_by.id == 0

    oso.load_str('allow("user", "read", post: Post) if post.id = 0;')

    # Baked queries disabled for sqlalchemy_oso.session.scoped_session.
    authorized_session = scoped_session(
        lambda: oso, lambda: "user", lambda: {Post: "read"}
    )
    authorized_session.configure(bind=engine)

    assert authorized_session.query(User).count() == 0

    authorized_posts = authorized_session.query(Post)
    assert authorized_posts.count() == 1
    first_authorized_post = authorized_posts[0]
    assert first_post.id == first_authorized_post.id

    # Should not be able to view the post's creator because there's no rule
    # permitting access to "read" users.
    assert first_authorized_post.created_by is None


def test_checked_permissions(engine, oso, fixture_data):
    """Test a query involving multiple models."""
    oso.load_str(
        """allow("user", "read", post: Post) if post.id = 1;
           allow("user", "view", user: User) if user.id = 0;
           allow("user", "view", user: User) if user.id = 1;
           allow("all_posts", "read", _: Post);"""
    )

    # Not applying any authorization to this session.
    session1 = AuthorizedSession(oso, "user", checked_permissions=None, bind=engine)
    posts1 = session1.query(Post)
    assert posts1.count() == 9
    assert posts1[0].created_by_id == 0
    assert posts1[0].created_by.username == "foo"
    users1 = session1.query(User)
    assert users1.count() == 3

    # Deny access to every model for this session by omission.
    session2 = AuthorizedSession(oso, "user", checked_permissions={}, bind=engine)
    posts2 = session2.query(Post)
    assert posts2.count() == 0
    users2 = session2.query(User)
    assert users2.count() == 0

    # Deny access to specific models for this session by inclusion.
    session3 = AuthorizedSession(
        oso, "user", checked_permissions={Post: None, User: None}, bind=engine
    )
    posts3 = session3.query(Post)
    assert posts3.count() == 0
    users3 = session3.query(User)
    assert users3.count() == 0

    # Allow access to one model but not the other for this session.
    session4 = AuthorizedSession(
        oso, "user", checked_permissions={Post: "read"}, bind=engine
    )
    posts4 = session4.query(Post)
    assert posts4.count() == 1
    assert posts4[0].id == 1
    assert posts4[0].created_by_id == 0
    assert posts4[0].created_by is None
    users4 = session4.query(User)
    assert users4.count() == 0

    # Allow access to multiple models with multiple actions for this session.
    session4 = AuthorizedSession(
        oso, "user", checked_permissions={Post: "read", User: "view"}, bind=engine
    )
    posts4 = session4.query(Post)
    assert posts4.count() == 1
    assert posts4[0].id == 1
    assert posts4[0].created_by_id == 0
    assert posts4[0].created_by.username == "foo"
    users4 = session4.query(User)
    assert users4.count() == 2


def test_register_models_declarative_base():
    """Test that `register_models()` registers models."""
    from oso import Oso
    from polar.exceptions import DuplicateClassAliasError

    from sqlalchemy_oso.auth import register_models

    from .models import Category, ModelBase, Tag

    oso = Oso()
    register_models(oso, ModelBase)

    for m in [Category, Post, Tag, User]:
        with pytest.raises(DuplicateClassAliasError):
            oso.register_class(m)


@pytest.mark.skipif(
    USING_SQLAlchemy_v1_3, reason="testing SQLAlchemy 1.4 functionality"
)
def test_register_models_registry():
    """Test that `register_models()` works with a SQLAlchemy 1.4-style
    registry."""
    # TODO(gj): remove type ignore once we upgrade to 1.4-aware MyPy types.
    from sqlalchemy.orm import registry  # type: ignore
    from sqlalchemy import Table, Column, Integer
    from oso import Oso
    from polar.exceptions import DuplicateClassAliasError

    from sqlalchemy_oso.auth import register_models

    mapper_registry = registry()

    user_table = Table(
        "user",
        mapper_registry.metadata,
        Column("id", Integer, primary_key=True),
    )

    class User:
        pass

    mapper_registry.map_imperatively(User, user_table)

    post_table = Table(
        "post",
        mapper_registry.metadata,
        Column("id", Integer, primary_key=True),
    )

    class Post:
        pass

    mapper_registry.map_imperatively(Post, post_table)

    oso = Oso()
    register_models(oso, mapper_registry)

    for m in [Post, User]:
        with pytest.raises(DuplicateClassAliasError):
            oso.register_class(m)
