"""Utility functions to support database interactions in tests."""

import collections
from contextlib import contextmanager
from functools import wraps
from unittest import mock
from sqlalchemy import inspect

from owsmysql import db


def seed_models(models, db_connection=None):
    """Save the given model(s) to the DB."""
    models = list(models)
    db_connection = db_connection or db.default
    assert db_connection
    assert isinstance(db_connection, db.TestConnection)

    with db_connection.session() as session:
        for model in models:
            session.add(model)
        session.flush()

        # detach the objects from this session so tests can interrogate them
        for model in models:
            session.expunge(model)


def create_schema():
    """Create all known tables definitions."""
    for db_connection in db.connections.__dict__.values():
        db_connection.create_all()


def drop_schema():
    """Drop all known tables definitions."""
    for db_connection in db.connections.__dict__.values():
        db_connection.drop_all()


def empty(function):
    """Test schema.

    Decorator that creates the test DB schema before a function call and
    tears the schema down after the function call has finished.

    This just creates the schema and does not seed data. Individual test cases
    can use factories to seed data as needed.

    Usage:

    @testdb.empty
    def test_smth()
        pass

    Args:
        Function (func): function to be called after creating the test schema.

    Returns:
        Function: The decorated function.

    """
    @wraps(function)
    def call_function_within_db_context(*args, **kwargs):
        create_schema()
        try:
            function_return = function(*args, **kwargs)
        finally:
            drop_schema()
        return function_return
    return call_function_within_db_context


def with_fixtures(*args, db_connection=None):
    """Wrap function to use DB filled with fixtures.

    Provide models or lists/tuples of models as arguments, e.g.
    fixture1 = model()
    fixture2 = model()
    fixture3 = model()
    fixture_set = (fixture2, fixture3)

    @testdb.with_fixtures(fixture1, fixture_set)
    def test_smth():
        pass
    """
    fixtures_list = []
    for arg in args:
        if isinstance(arg, collections.abc.Sequence):
            fixtures_list.extend(arg)
        else:
            fixtures_list.append(arg)

    def wrapper_factory(function):
        @wraps(function)
        def wrapper(*args, **kwargs):
            seed_models(fixtures_list, db_connection=db_connection)
            try:
                return function(*args, **kwargs)
            finally:
                for model in fixtures_list:
                    state = inspect(model)
                    state.key = None
        return empty(wrapper)

    return wrapper_factory


def mock_db_session(mocker):
    """Create a mock database session.

    Also mock the db_session context manager to use the mock session.
    """
    mock_session = mock.Mock(query=mock.Mock())

    @contextmanager
    def fake_session_manager(*args):
        yield mock_session

    mocker.patch.object(
        db.MySQLConnection, 'session', fake_session_manager)
    return mock_session