#!/usr/bin/env python3

import argparse
import glob
import os
import subprocess
import sys
import threading

from contextlib import contextmanager

CONFIG = {}
DC_ENV = {}


def read_proc_callback(proc, output_handler=lambda line: line):
    for line in iter(proc.stdout.readline, b''):
        output_handler(line)


def read_config(fpath='../.env.dev'):
    with open(fpath, 'r') as f:
        content = f.read()
    config = {}
    for line in content.splitlines():
        line = line.strip()
        if line.startswith('#'):
            continue
        key, val = line.split('=', 1)
        config[key] = val
    return config


def is_container_running(container):
    return bool(subprocess.check_output(['docker', 'ps', '-q', '--filter', f'name={container}'], env=DC_ENV))


def get_containers_ids(*additional_docker_ps_args):
    names = subprocess.check_output(['docker', 'ps', '-q', *additional_docker_ps_args], env=DC_ENV)
    return names.decode('utf-8').split()


def stop(containers=None):
    """ Stop passed containers or all containers if no names passed
        :param containers: list|tuple names or string name or None, meaning all
        :return: list of stopped containers names"""
    if containers is None:
        containers = get_containers_ids()
    if not isinstance(containers, (list, tuple)):
        containers = [containers]
    if containers:
        arguments = ['docker', 'stop', *containers]
        print(f'Stopping containers {containers}')
        proc = subprocess.run(arguments, env=DC_ENV)
    return containers


def down():
    """ Stop and remove containers """
    arguments = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/selenium.yml',
                 '-f', 'compose/flower.yml', '-f', 'compose/pgadmin.yml',
                 '-f', 'compose/el.yml', '-f', 'compose/correlator.yml',
                 '--env-file=../.env.dev', '-f', 'compose/kibana.yml',
                 '-f', 'compose/license.yml', 'down']
    print('Docker-compose down')
    subprocess.run(arguments, env=DC_ENV)


def exec_django_container_cmd(*cmds):
    global CONFIG
    subprocess.run(['docker', 'exec', '-ti', CONFIG.get('DOCKER_DJANGO_NAME'), *cmds], env=DC_ENV)


def wait_process_stop(proc, timeout=None):
    """ Wait for process which supports interruption, i.e. it can stop itself on interruption """
    try:
        try:
            return proc.wait(timeout=timeout)
        except KeyboardInterrupt:
            print('Stopping process gracefully, repeat interruption to stop immediately')
            proc.wait()
    except:
        proc.kill()
        proc.wait()
        raise


def wait_till_output_bytes_line(proc,
                                bytes_lines_to_wait=[b'Starting Gnuicorn server', b'Starting Django dev server'],
                                print_output=True,
                                reraise=False):
    if not isinstance(bytes_lines_to_wait, (tuple, list)):
        bytes_lines_to_wait = [bytes_lines_to_wait]
    try:
        print('Waiting for docker ready')
        for line in iter(proc.stdout.readline, b''):
            if print_output:
                print(line.decode('utf-8')[:-1])
            if any((l in line for l in bytes_lines_to_wait)):
                return True
        return False
    except KeyboardInterrupt:
        print('Stopping process gracefully, repeat interruption to stop immediately')
        return False


@contextmanager
def docker_on(use_gunicorn=False, use_selenium=False, use_el=False, use_correlator=False, use_kibana=False,
              use_flower=False, use_pgadmin=False, wait_till_loaded=True):
    global CONFIG
    arguments = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml']
    was_running = []

    needed_conts = ('DOCKER_DJANGO_NAME',)
    was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts])

    if use_selenium:
        needed_conts = ('DOCKER_FIREFOX_NODE_NAME', 'DOCKER_CHROME_NODE_NAME', 'DOCKER_SELENIUM_NAME')
        was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts])
        arguments.append('-f')
        arguments.append('compose/selenium.yml')

    if use_flower:
        needed_conts = ('DOCKER_FLOWER_NAME',)
        was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts])
        arguments.append('-f')
        arguments.append('compose/flower.yml')

    if use_pgadmin:
        needed_conts = ('DOCKER_PGADMIN_SERVER',)
        was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts])
        arguments.append('-f')
        arguments.append('compose/pgadmin.yml')

    if use_el:
        needed_conts = ('DOCKER_ELASTIC_NAME', 'DOCKER_VECTOR_NAME')
        was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts])
        arguments.append('-f')
        arguments.append('compose/el.yml')
        arguments.append('--env-file=../.env.dev')

    if use_correlator:
        arguments.append('-f')
        arguments.append('compose/correlator.yml')

    if use_kibana:
        needed_conts = ('DOCKER_KIBANA_NAME',)
        was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts])
        arguments.append('-f')
        arguments.append('compose/kibana.yml')

    if not use_gunicorn:
        DC_ENV['SERVE'] = 'django'

    was_running = all([was_running, *was_running])
    arguments.append('up')

    proc = None
    if not was_running:
        print(f'Starting containers')
        kwargs = {}
        if wait_till_loaded:
            kwargs = dict(stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        proc = subprocess.Popen(arguments, env=DC_ENV, **kwargs)
        if wait_till_loaded:
            if not wait_till_output_bytes_line(proc):
                stop()
                raise KeyboardInterrupt
    try:
        if proc and wait_till_loaded:  # prevent process queue overflow
            def drop_lines(line): pass

            threading.Thread(target=read_proc_callback, args=(proc, drop_lines)).start()
        yield proc
    except:
        if proc:
            proc.kill()
            proc.wait()
            raise
    finally:
        if not was_running:
            stop()
        if proc:
            wait_process_stop(proc)


def test(enable_warning=False, add_args=[]):
    """ Call process as subprocess.call, but on interruption wait for termination of process """
    global CONFIG
    test_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f',
                 'compose/selenium.yml', 'exec', CONFIG.get('DOCKER_DJANGO_NAME'), 'pytest']

    if not enable_warning:
        test_args.append('--disable-warnings')

    test_args.extend(add_args)

    with docker_on(use_selenium=True, wait_till_loaded=True):
        print('Starting tests')
        test_proc = subprocess.Popen(
            test_args, env=DC_ENV, stdout=subprocess.STDOUT, stderr=subprocess.STDOUT)
        wait_process_stop(test_proc)
        sys.exit(test_proc.returncode)


def ci_test(add_args):
    """ Used to run tests in CI """
    global CONFIG
    test_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f',
                 'compose/selenium.yml', 'exec', '-T', CONFIG.get('DOCKER_DJANGO_NAME'), 'pytest']
    test_args.extend(add_args)

    with docker_on(use_selenium=True, wait_till_loaded=True):
        print('Starting tests')
        test_proc = subprocess.Popen(test_args, env=DC_ENV)
        wait_process_stop(test_proc)

        print("Exit code:", test_proc.returncode)
        sys.exit(test_proc.returncode)


def ci_test_coverage(add_args):
    """ Used to check test coverage"""
    global CONFIG
    DC_ENV["COVERAGE_FILE"] = "public/.coverage"

    test_args = ['docker-compose',
                 '-f',
                 'compose/django.yml',
                 '-f',
                 'compose/license.yml',
                 '-f',
                 'compose/selenium.yml',
                 'exec',
                 '-T',
                 CONFIG.get('DOCKER_DJANGO_NAME'),
                 'coverage',
                 "run",
                 "-m",
                 "pytest",
                 "-m",
                 "'merge or selenium'"]
    report_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', 'exec', '-T',
                   CONFIG.get('DOCKER_DJANGO_NAME'), 'coverage', "html", "-d", "public/test_coverage/"]

    test_args.extend(add_args)

    result = 0

    with docker_on(use_selenium=True, wait_till_loaded=True):
        print(f"Starting tests {test_args}")
        test_proc = subprocess.Popen(test_args, env=DC_ENV)
        wait_process_stop(test_proc)

        result = test_proc.returncode

    with docker_on(use_selenium=True, wait_till_loaded=True):
        print(f"Processing report {report_args}")
        test_report = subprocess.Popen(report_args, env=DC_ENV)
        wait_process_stop(test_report)

    sys.exit(result)


def parse_args(argv=sys.argv[1:]):
    parser = argparse.ArgumentParser(
        description=f'Project starter. To open web ui http://localhost:{CONFIG.get("WEB_UI_PORT", 9090)}/ with admin:nimda')
    sub_parsers = parser.add_subparsers(title='Commands', dest='cmd', required=True)

    run_parser = sub_parsers.add_parser(name='run', help='Run custom configuration')
    run_parser.add_argument('--with-selenium', dest='use_selenium', action='store_true', help='Add selenium')
    run_parser.add_argument('--with-flower', dest='use_flower', action='store_true',
                            help='Add flower. Use http://localhost:5555')
    run_parser.add_argument('--with-pgadmin', dest='use_pgadmin', action='store_true',
                            help='Add Pgadmin. Use http://localhost:5050 with pgadmin4@pgadmin.org : admin')
    run_parser.add_argument('--with-ev', dest='use_ev', action='store_true', help='Add Vector and Elasticsearch')
    run_parser.add_argument('--with-cev', dest='use_cev', action='store_true',
                            help='Add Vector, Elasticsearch and correlator')
    run_parser.add_argument('--with-kibana', dest='use_kibana', action='store_true',
                            help='Add Kibana. Use http://localhost:5601 with elastic:changeme')
    run_parser.add_argument('--with-gunicorn', dest='use_gunicorn', action='store_true',
                            help='Use gunicorn server instead of django')
    run_parser.add_argument('--with-evk', dest='use_evk', action='store_true',
                            help='Add Vector, Elasticsearch and Kibana')
    run_parser.add_argument('--with-cevk', dest='use_cevk', action='store_true',
                            help='Add Vector, Elasticsearch, correlator and Kibana')

    test_parser = sub_parsers.add_parser(name='test', aliases=['tests'], help='Run tests')
    test_parser.add_argument('--enable-warning', dest='enable_warning', action='store_true',
                             help='Enable warning for tests')

    ci_test_parser = sub_parsers.add_parser(name="ci_test",
                                            help="Run tests in CI. Differs from the test function in the way of running docker compose. Use for CI only.")

    ci_test_coverage = sub_parsers.add_parser(name="ci_test_coverage",
                                              help="Run test coverage in CI. Differs from the test function in the way of running docker compose. Use for CI only.")

    build_parser = sub_parsers.add_parser(name='build', aliases=['rebuild'], help='Rebuild containers')
    build_parser.add_argument('--no-cache', dest='no_cache', action='store_true', help='Rebuild without cache')

    stop_parser = sub_parsers.add_parser(name='stop', help='Stop containers')
    down_parser = sub_parsers.add_parser(name='down', help='Stop containers and remove containers')

    django_restart_db_parser = sub_parsers.add_parser(name='restart-django', aliases=['reload-django'],
                                                      help='Restart django container')
    sh_parser = sub_parsers.add_parser(name='sh', help='Run shell in django container')
    psh_parser = sub_parsers.add_parser(name='psh', help='Run python django shell in django container')
    psql_parser = sub_parsers.add_parser(name='psql', help='Run postgresql shell')
    clear_db_parser = sub_parsers.add_parser(name='clear-db', aliases=['db-clear'],
                                             help="Remove migrations which are not in git and drop container's data and tables from database")
    clear_db_parser.add_argument('-n', '--dry-run', dest='dry_run', action='store_true',
                                 help='Show actions without actual deletion of info')
    webui_parser = sub_parsers.add_parser(name='ui', help='Open web UI')
    webpdb_parser = sub_parsers.add_parser(name='wpdb', help='Open web PDB debugger')

    args, unknownargs = parser.parse_known_args(argv)
    return args, unknownargs

def main():
    global CONFIG, DC_ENV
    CONFIG.update(read_config())
    DC_ENV = {**os.environ.copy(), **CONFIG}

    args, unknownargs = parse_args()

    if args.cmd == 'run':
        if args.use_cevk:
            args.use_evk, args.user_cev = True, True
        if args.use_cev:
            args.use_ev = True
        if args.use_evk:
            args.use_ev, args.use_kibana = True, True
        with docker_on(args.use_gunicorn, args.use_selenium, args.use_ev, args.use_cevk, args.use_kibana,
                       args.use_flower, args.use_pgadmin, wait_till_loaded=False) as proc:
            if proc:
                wait_process_stop(proc)

    elif args.cmd in ('test', 'tests'):
        test(args.enable_warning, unknownargs)

    elif args.cmd in ('ci_test'):
        ci_test(unknownargs)
    elif args.cmd in ("ci_test_coverage"):
        ci_test_coverage(unknownargs)

    elif args.cmd in ('build', 'rebuild'):
        stop()
        arguments = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f',
                     'compose/selenium.yml', '-f',
                     'compose/pgadmin.yml', '--env-file=../.env.dev', '-f', 'compose/el.yml', '-f',
                     'compose/correlator.yml', '-f', 'compose/kibana.yml', 'build']

        if args.no_cache:
            arguments.append('--no-cache')

        proc = subprocess.run(arguments, env=DC_ENV)

    elif args.cmd == 'stop':
        stop()

    elif args.cmd == 'down':
        down()

    elif args.cmd in ('restart-django', 'reload-django'):
        subprocess.run(['docker', 'restart', CONFIG.get('DOCKER_DJANGO_NAME')], env=DC_ENV)

    elif args.cmd in ('sh',):
        with docker_on(wait_till_loaded=True) as proc:
            exec_django_container_cmd('bash', *unknownargs)
    elif args.cmd in ('psh',):
        with docker_on(wait_till_loaded=True) as proc:
            exec_django_container_cmd(*['bash', '-c', 'python3 manage.py shell'], *unknownargs)
    elif args.cmd in ('psql',):
        with docker_on(wait_till_loaded=True) as proc:
            arguments = ['psql', '-U', CONFIG.get('POSTGRES_USER'), '-d', CONFIG.get("POSTGRES_DB")]
            subprocess.run(['docker', 'exec', '-ti', CONFIG.get('DOCKER_DB_NAME'), *arguments], env=DC_ENV)
    elif args.cmd in ('clear-db', 'db-clear'):
        print('Clearing migrations')
        git_clean_opts = '-xf' + (args.dry_run and 'n' or '')
        subprocess.run(['sudo', 'git', 'clean', git_clean_opts, *glob.glob("../*/migrations/")], env=DC_ENV)
        subprocess.run(['sudo', 'git', 'clean', git_clean_opts, 'compose/config/elk/vector/pipeline/armaif_*.toml'],
                       env=DC_ENV)
        subprocess.run(['sudo', 'git', 'clean', git_clean_opts, 'compose/config/elk/vector/pipeline/endpoint_*.toml'],
                       env=DC_ENV)
        print("Applying docker-compose down to all containers")
        if not args.dry_run:
            down()
    elif args.cmd == 'ui':
        import webbrowser
        webbrowser.open(f'http://localhost:{CONFIG.get("WEB_UI_PORT")}')
    elif args.cmd == 'wpdb':
        import webbrowser
        webbrowser.open(f'http://localhost:{CONFIG.get("WEB_PDB_PORT")}')


if __name__ == '__main__':
    main()
