#!/usr/bin/env python3

from contextlib import contextmanager
import argparse
import subprocess
from typing import List, Dict
import sys
import yaml
import os
from tempfile import NamedTemporaryFile


script_dir = os.path.dirname(os.path.realpath(__file__))

inside_docker_db_facts_local_path =\
    f'{script_dir}/inside-docker-dbfacts.yml'

itest_volumes = {
    inside_docker_db_facts_local_path: '/root/.dbfacts.yml',
}


@contextmanager
def dockerized_dbs():
    docker_compose_start()
    try:
        yield
    finally:
        docker_compose(["stop"])


@contextmanager
def set_s3_scratch_bucket():
    # https://stackoverflow.com/questions/2059482/python-temporarily-modify-the-current-processs-environment
    _environ = dict(os.environ)
    os.environ['SCRATCH_S3_URL'] = 's3://circleci-scratch/'
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(_environ)


@contextmanager
def set_gcs_scratch_bucket():
    # https://stackoverflow.com/questions/2059482/python-temporarily-modify-the-current-processs-environment
    _environ = dict(os.environ)
    os.environ['SCRATCH_GCS_URL'] = 'gs://bluelabs-test-recordsmover/'
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(_environ)


@contextmanager
def local_dockerized_dbfacts():
    """
    To connect to docker-compose from outside of Docker, we need to
    ask docker-compose what port it mapped things to.
    """
    # https://stackoverflow.com/questions/2059482/python-temporarily-modify-the-current-processs-environment
    _environ = dict(os.environ)
    try:
        with NamedTemporaryFile() as f:
            vertica_port =\
                subprocess.check_output([
                    './itest-dc', 'port', 'verticadb', '5433'
                ]).decode('utf8').rstrip().split(':')[1]
            postgres_port =\
                subprocess.check_output([
                    './itest-dc', 'port', 'postgresdb', '5432'
                ]).decode('utf8').rstrip().split(':')[1]
            mysql_port =\
                subprocess.check_output([
                    './itest-dc', 'port', 'mysqldb', '3306'
                ]).decode('utf8').rstrip().split(':')[1]
            db_facts = {
                'dbs': {
                    'dockerized-vertica': {
                        'exports': {
                            'host': 'localhost',
                            'port': vertica_port,
                            'database': 'docker',
                            'type': 'vertica',
                            'protocol': 'vertica',
                            'user': 'dbadmin',
                            'password': "",
                        }
                    },
                    'dockerized-postgres': {
                        'exports': {
                            'host': 'localhost',
                            'port': postgres_port,
                            'database': 'postgres',
                            'type': 'postgres',
                            'protocol': 'postgres',
                            'user': 'postgres',
                            'password': 'hunter2',
                        }
                    },
                    'dockerized-mysql': {
                        'exports': {
                            # This needs to be 127.0.0.1, because if
                            # this is the string 'localhost', the
                            # MySQL driver wants to use Unix domain
                            # sockets to connect, which won't work
                            # because this is a tunnelled port.
                            'host': '127.0.0.1',
                            'port': mysql_port,
                            'database': 'mysqlitest',
                            'type': 'mysql',
                            'protocol': 'mysql',
                            'user': 'mysqluser',
                            'password': 'hunter2',
                        }
                    },
                }
            }
            yaml_output = yaml.dump(db_facts)
            f.write(yaml_output.encode('utf8'))
            f.flush()
            os.environ['DB_FACTS_PATH'] = f.name
            yield
    finally:
        os.environ.clear()
        os.environ.update(_environ)


def docker_build() -> None:
    subprocess.check_call("make docker",
                          shell=True,
                          cwd=f"{script_dir}/../..")


def docker_compose(args: List[str],
                   prefixes: List[str] = []) -> None:
    subprocess.check_call(prefixes + [f"{script_dir}/itest-dc"] + args)


def docker_compose_run(cmd: List[str],
                       volumes: Dict[str, str] = itest_volumes,
                       prefixes: List[str] = [],
                       cwd: str = "/usr/src/app") -> None:
    volumes_args = []
    for local_path, docker_path in volumes.items():
        volumes_args.append('-v')
        volumes_args.append(f"{local_path}:{docker_path}")
    docker_compose(["run", "--rm", "--workdir", cwd] + volumes_args + ["records_mover"] + cmd,
                   prefixes=prefixes)


def docker_compose_shell() -> None:
    docker_compose_run(["/bin/sh"])


def docker_compose_start() -> None:
    print("Running docker_compose start verticadb postgresdb mysqldb...", file=sys.stderr)
    docker_compose(["up", "--no-start"])
    docker_compose(["start", "verticadb", "postgresdb", "mysqldb"])
    docker_compose_run(['tests/integration/wait-for-vertica.sh'])
    print("Verified Vertica is up and listening", file=sys.stderr)
    docker_compose_run(['tests/integration/wait-for-postgres.sh'])
    print("Verified Postgres is up and listening", file=sys.stderr)
    docker_compose_run(['tests/integration/wait-for-mysql.sh'])
    print("Verified MySQL is up and listening", file=sys.stderr)


def run_test_build(args):
    if (args.docker):
        docker_build()
    else:
        print(f"Nothing to build (did you intend to use --docker?)", file=sys.stderr)


def run_test_up(args):
    if (args.docker):
        docker_compose(["up", "--no-start"])
    else:
        print(f"Nothing to bring up (did you intend to use --docker?)", file=sys.stderr)


def run_test_down(args):
    if (args.docker):
        docker_compose(["down"])
    else:
        print(f"Nothing to bring down (did you intend to use --docker?)", file=sys.stderr)


def run_test_ps(args):
    if (args.docker):
        docker_compose(["ps"])
    else:
        print(f"Nothing to list (did you intend to use --docker?)", file=sys.stderr)


def run_test_shell(args):
    if (args.docker):
        docker_compose_shell()
    else:
        with dockerized_dbs(), local_dockerized_dbfacts():
            subprocess.check_call(["/bin/bash"])


def run_test_cli(args):
    with dockerized_dbs(), set_s3_scratch_bucket(), set_gcs_scratch_bucket():
        if (args.docker):
            docker_compose_run(["./cli-tests.sh"],
                               prefixes=["with-aws-creds", "circleci"],
                               cwd="/usr/src/app/tests/integration/cli")
        else:
            with local_dockerized_dbfacts():
                subprocess.check_call(["with-aws-creds", "circleci",
                                       "./cli-tests.sh"],
                                      cwd="cli")


def run_test_vertica_no_s3(args):
    with dockerized_dbs():
        if (args.docker):
            docker_compose_run(['with-db', 'dockerized-vertica',
                                'pytest', '--xunit-file=pytest.xml', '.'],
                               cwd="/usr/src/app/tests/integration/records/single_db")

        else:
            with local_dockerized_dbfacts():
                subprocess.check_call(['with-db', 'dockerized-vertica',
                                       'pytest', '--xunit-file=pytest.xml', '.'],
                                      cwd=f"{script_dir}/records/single_db")


def run_test_vertica_s3(args):
    with dockerized_dbs(), set_s3_scratch_bucket():
        if (args.docker):
            docker_compose_run(['with-db', 'dockerized-vertica',
                                'pytest', '--xunit-file=pytest.xml', '.'],
                               prefixes=["with-aws-creds", "circleci"],
                               cwd="/usr/src/app/tests/integration/records/single_db")
        else:
            with local_dockerized_dbfacts():
                subprocess.check_call(["with-db", "dockerized-vertica",
                                       "with-aws-creds", "circleci",
                                       "pytest", "--xunit-file=pytest.xml", "."],
                                      cwd=f"{script_dir}/records/single_db")


def run_test_mysql(args):
    with dockerized_dbs():
        if (args.docker):
            docker_compose_run(['with-db', 'dockerized-mysql',
                                'pytest', '--xunit-file=pytest.xml', '.'],
                               prefixes=["with-aws-creds", "circleci"],
                               cwd="/usr/src/app/tests/integration/records/single_db")
        else:
            with local_dockerized_dbfacts():
                subprocess.check_call(["with-db", "dockerized-mysql",
                                       "pytest", "--xunit-file=pytest.xml", "."],
                                      cwd=f"{script_dir}/records/single_db")


def run_test_postgres(args):
    with dockerized_dbs():
        if (args.docker):
            docker_compose_run(['with-db', 'dockerized-postgres',
                                'pytest', '--xunit-file=pytest.xml', '.'],
                               prefixes=["with-aws-creds", "circleci"],
                               cwd="/usr/src/app/tests/integration/records/single_db")
        else:
            with local_dockerized_dbfacts():
                subprocess.check_call(["with-db", "dockerized-postgres",
                                       "pytest", "--xunit-file=pytest.xml", "."],
                                      cwd=f"{script_dir}/records/single_db")


def run_test_redshift_s3(args):
    with set_s3_scratch_bucket():
        if (args.docker):
            docker_compose_run(['with-db', 'demo-itest',
                                'pytest', '--xunit-file=pytest.xml', '.'],
                               prefixes=["with-aws-creds", "circleci"],
                               cwd="/usr/src/app/tests/integration/records/single_db")
        else:
            subprocess.check_call(['with-db', 'demo-itest',
                                   "with-aws-creds", "circleci",
                                   'pytest', '--xunit-file=pytest.xml', '.'],
                                  cwd=f"{script_dir}/records/single_db")


def run_test_redshift_no_s3(args):
    if (args.docker):
        docker_compose_run(['with-db', 'demo-itest',
                            'pytest', '--xunit-file=pytest.xml', '.'],
                           prefixes=["with-aws-creds", "circleci"],
                           cwd="/usr/src/app/tests/integration/records/single_db")
    else:
        subprocess.check_call(['with-db', 'demo-itest',
                               'pytest', '--xunit-file=pytest.xml', '.'],
                              cwd=f"{script_dir}/records/single_db")


def run_test_bigquery_no_gcs(args):
    if (args.docker):
        docker_compose_run(['with-db', 'bltoolsdevbq-bq_itest',
                            'pytest', '--xunit-file=pytest.xml', '.'],
                           cwd="/usr/src/app/tests/integration/records/single_db")
    else:
        subprocess.check_call(['with-db', 'bltoolsdevbq-bq_itest',
                               'pytest', '--xunit-file=pytest.xml', '.'],
                              cwd=f"{script_dir}/records/single_db")


def run_test_bigquery_gcs(args):
    with set_gcs_scratch_bucket():
        if (args.docker):
            docker_compose_run(['with-db', 'bltoolsdevbq-bq_itest',
                                'pytest', '--xunit-file=pytest.xml', '.'],
                               cwd="/usr/src/app/tests/integration/records/single_db")
        else:
            subprocess.check_call(['with-db', 'bltoolsdevbq-bq_itest',
                                   'pytest', '--xunit-file=pytest.xml', '.'],
                                  cwd=f"{script_dir}/records/single_db")


def run_test_table2table(args):
    with set_s3_scratch_bucket(), set_gcs_scratch_bucket(), dockerized_dbs():
        if (args.docker):
            docker_compose_run(['python3', '-m',
                                'records.multi_db.test_records_table2table'],
                               prefixes=["with-aws-creds", "circleci"],
                               cwd="/usr/src/app/tests/integration")
        else:
            with local_dockerized_dbfacts():
                subprocess.check_call(["with-aws-creds", "circleci",
                                       'python3', '-m',
                                       'records.multi_db.test_records_table2table'],
                                      cwd=".")


def run_test(args, target, parser):
    print(f"Executing {target} test", file=sys.stderr)

    run_test_target_to_func = {
        'build': run_test_build,
        'up': run_test_up,
        'down': run_test_down,
        'ps': run_test_ps,
        'shell': run_test_shell,
        'cli': run_test_cli,
        'vertica-no-s3': run_test_vertica_no_s3,
        'vertica-s3': run_test_vertica_s3,
        'mysql': run_test_mysql,
        'postgres': run_test_postgres,
        'redshift-s3': run_test_redshift_s3,
        'redshift-no-s3': run_test_redshift_no_s3,
        'bigquery-no-gcs': run_test_bigquery_no_gcs,
        'bigquery-gcs': run_test_bigquery_gcs,
        'table2table': run_test_table2table,
    }

    func = run_test_target_to_func.get(target)
    if func is not None:
        func(args)
    else:
        parser.print_help()
        sys.exit(1)


def main():
    tests = {
        'cli': 'Run bash-based multi-source/target copy tests',
        'mysql': 'Run load/unload suite against Dockerized MySQL',
        'postgres': 'Run load/unload suite against Dockerized PostgreSQL',
        'vertica-s3': 'Run load/unload suite against Dockerized Vertica, using S3',
        'vertica-no-s3': 'Run load/unload suite against Dockerized Vertica, using streams',
        'redshift-s3': 'Run load/unload suite against BlueLabs Demo Redshift, using S3',
        'redshift-no-s3': 'Run load/unload suite against BlueLabs Demo Redshift, using '
        'INSERT/SELECT',
        'bigquery-no-gcs': 'Run load/unload suite against BlueLabs BigQuery, using streams',
        'bigquery-gcs': 'Run load/unload suite against BlueLabs BigQuery, using GCS',
        'table2table': 'Run Python-based cross-database copy tests',
    }

    other_targets = {
        'all': 'Run all of the above tests',
        'build': 'Rebuild the Docker image used in --docker testing',
        'up': 'Bring up the docker-compose environment',
        'down': 'Bring down the docker-compose environment',
        'shell': 'Run a shell environment inside the docker-compose records_mover container',
        'ps': 'Show currently running docker-compose containers',
    }

    epilog = "The following targets are supported:\n\n" + "\n\n".join([
        f"  {key}: {value}"
        for key, value in
        {**tests, **other_targets}.items()
    ])
    parser = argparse.ArgumentParser(description='Run records-mover integration tests',
                                     formatter_class=argparse.RawTextHelpFormatter,
                                     epilog=epilog)
    parser.add_argument('--docker', action="store_true", default=False,
                        help='Run mover itself from within Docker')

    parser.add_argument('targets',
                        nargs='+',
                        choices=list(tests.keys()) + list(other_targets.keys()),
                        help='Action to take/test to run')

    args = parser.parse_args()

    # Unset any existing scratch bucket, as it should be set
    # explicitly here or not used entirely, depending on the test:
    if 'SCRATCH_S3_URL' in os.environ:
        del os.environ['SCRATCH_S3_URL']

    if 'SCRATCH_GCS_URL' in os.environ:
        del os.environ['SCRATCH_GCS_URL']

    targets_lists = [
        [target] if target != 'all'
        else tests.keys()
        for target in args.targets
    ]

    def flatten(list_of_lists):
        return [
            item
            for items in list_of_lists
            for item in items
        ]

    targets = flatten(targets_lists)

    print(f"Integration test targets: {' '.join(targets)}", file=sys.stderr)

    for target in targets:
        run_test(args, target, parser)


if __name__ == "__main__":
    main()
