"""
Tasks for working with a Docker Compose setup for testing and development.

:copyright: Copyright 2013-2016, Dorian Pula <dorian.pula@amber-penguin-software.ca>
:license: AGPL v3+
"""
import contextlib
import os
import sys
import traceback

import arrow
import invoke as inv
import sqlalchemy
from sqlalchemy import exc as sql_error
import sqlalchemy_utils

from rookeries import database as rookeries_db
from tasks.utils import environment
from tests.fixtures import sample_site


DEFAULT_DB_CONNECTION = 'postgresql+psycopg2://admin:password@localhost:5432/rookeries'


def get_db_connection_uri(db_connection, db_name):
    db_connection = os.environ.get(rookeries_db.EXTERNAL_DB_URI_PARAM, db_connection)
    return '/'.join((db_connection.rsplit('/', maxsplit=1)[0], db_name))


@inv.task
def bootstrap(ctx, db_connection=DEFAULT_DB_CONNECTION, db_name='rookeries', reset_db=False):
    """
    Bootstrap the database with sample code.

    :param ctx: Context of the invoke task.
    :param db_connection: The database connection string.  Gets overridden with the environment variable
        ROOKERIES_SQLALCHEMY_DATABASE_URI.
    :param db_name: The name of the database to create.
    :param reset_db: Flag to reset a database and tables if they exist.
    """

    db_connection = get_db_connection_uri(db_connection, db_name)
    if not sqlalchemy_utils.database_exists(db_connection):
        sqlalchemy_utils.create_database(db_connection)

    if sqlalchemy_utils.database_exists(db_connection):
        if not reset_db:
            return
        sqlalchemy_utils.drop_database(db_connection)

    with test_database(db_connection=db_connection, keep_test_db=True, test_db_name=db_name):
        pass


@contextlib.contextmanager
def test_database(db_connection=DEFAULT_DB_CONNECTION, keep_test_db=False,
                  test_db_name='rookeries_{}'.format(arrow.now().format('YYYY_MM_DD_HHmmss'))):
    """
    Context for creating and managing a test database.

    :param db_connection: The database connection string.  Gets overridden with the environment variable
        ROOKERIES_SQLALCHEMY_DATABASE_URI.
    :param keep_test_db: Keep the test data and database around.
    :param test_db_name: The name of the test database.
    """

    try:
        db_connection = os.environ.get(rookeries_db.EXTERNAL_DB_URI_PARAM, db_connection)
        db_connection = '/'.join((db_connection.rsplit('/', maxsplit=1)[0], test_db_name))

        if not sqlalchemy_utils.database_exists(db_connection):
            sqlalchemy_utils.create_database(db_connection)

        db_engine = sqlalchemy.create_engine(db_connection, echo=True)
        rookeries_db.ModelBase.metadata.create_all(db_engine)

        sample_site.populate_database(db_engine)
        modified_environ = {
            'ROOKERIES_SQLALCHEMY_DATABASE_URI': db_connection,
        }

        with environment(modified_environ):
            pass
            yield

    except sql_error.OperationalError as err:
        print(f'Unable to connect to the database because of "{err}" \n Stack: \n')
        traceback.print_tb(sys.exc_info()[2])
        raise err

    except Exception as err:
        print(f'Unable to run server test because of "{err}" \n Stack: \n')
        traceback.print_tb(sys.exc_info()[2])
        raise err

    finally:
        if not keep_test_db:
            sqlalchemy_utils.drop_database(db_connection)
