#!/usr/bin/env python3
import subprocess
import sys

import argparse

from enum import Enum, auto


YANDEX_ENVIRONMENT_TYPE_FILE_PATH = "/etc/yandex/environment.type"
REPLICATION_PASSWORD_FILE_PATH = "/etc/cauth-ldap-slave-slapd-config/replication_password"
REPLICATION_CONF_TEMPLATE_FILE_PATH = "/etc/cauth-ldap-slave-slapd-config/replication-settings.conf.template"
REPLICATION_CONF_FILE_PATH = "/etc/cauth-ldap-slave-slapd-config/replication-settings.conf"
OPENLDAP_USER = "openldap"
OPENLDAP_GROUP = "openldap"
SLAPD_NUMBER_OF_INSTANCES = 12  # TODO: Read file '/etc/default/slapd-instances', parse 'num_instances="12"'


class YandexEnvironment(Enum):
    UNKNOWN = auto()
    TESTING = auto()
    PRODUCTION = auto()


class YandexEnvironmentKnownTypes(str, Enum):
    PRODUCTION: str = "production"
    TESTING: str = "testing"


def parse_args() -> argparse.Namespace:
    """ Parses CLI arguments. """
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="cli_action")

    parser_get_ldap_replication_password = subparsers.add_parser("get-ldap-replication-password")
    parser_get_ldap_replication_password.add_argument("--yandex-vault-username", type=str, required=True)

    parser_setup_ldap_replication = subparsers.add_parser("setup-ldap-replication")
    parser_setup_ldap_replication.add_argument("--replication-master-host", type=str, required=True)

    parser_copy_database = subparsers.add_parser("copy-database")
    parser_copy_database.add_argument("--database-source-host", type=str, required=True)

    return parser.parse_args()


def copy_database(database_source_host: str):
    local_database_dirs = []
    for instance_number in range(1, _get_num_instances_configured()+1):
        local_database_dirs.append(f"/var/lib/slapd-instances/instance-{instance_number:02d}")

    remote_database_dir = "/var/lib/slapd-instances/instance-01"
    database_file_name = "data.mdb"

    print(f"Copying database from {database_source_host} to /tmp/{database_source_host}/{database_file_name}")
    _run_subprocess(
        (
            f"mkdir -p /tmp/{database_source_host} && "
            f"scp -r {database_source_host}:{remote_database_dir}/{database_file_name} "
            f"/tmp/{database_source_host}/{database_file_name}"
        ),
        shell=True,
    )

    for dst in local_database_dirs:
        print(f"Copying database from /tmp/{database_source_host}/{database_file_name} to {dst}")
        _run_subprocess(
            (
                f"mkdir -p {dst} && "
                f"cp -r /tmp/{database_source_host}/{database_file_name} {dst}/{database_file_name}"
            ),
            shell=True
        )


def get_replication_password_from_vault(yandex_vault_username: str):
    """ Runs 'yav-deploy'. """
    if not yandex_vault_username:
        raise Exception(f"Empty Yandex Vault username.")

    yandex_environment = _get_yandex_environment()
    config_path = {
        YandexEnvironment.PRODUCTION: "/etc/yandex/yav-deploy/cauth-ldap-slave-slapd-config-production.conf",
    }.get(yandex_environment)
    if not config_path:
        raise Exception(f"Environment '{yandex_environment}' not supported.")

    command = [
        "yav-deploy",
        "--file",
        config_path,
        "--rsa-login",
        yandex_vault_username,
        "--debug"
    ]

    print("Calling 'yav-deploy'...")
    result = _run_subprocess(command)

    if not result:
        print(f"""
* Check Yandex Vault username.
* Ensure that 'ForwardAgent Yes' is in your ~/.ssh/config file.
* Ensure that you have access to secret described in {config_path} file through ABC role.
""")
        sys.exit(1)


def setup_ldap_replication(replication_master_host: str):
    try:
        with open(REPLICATION_PASSWORD_FILE_PATH) as f:
            replication_password = f.read().rstrip()
    except FileNotFoundError:
        print(
            (
                f"Password file '{REPLICATION_PASSWORD_FILE_PATH}' does not exist, "
                f"run 'cauth-ldap-replication-manager get-ldap-replication-password'"
            )
        )
        sys.exit(1)

    _run_subprocess(
        (
            f"envsubst < "
            f"{REPLICATION_CONF_TEMPLATE_FILE_PATH} > "
            f" {REPLICATION_CONF_FILE_PATH}"
        ),
        shell=True,
        env=dict(
            LDAP_REPLICATION_MASTER_HOST=replication_master_host,
            LDAP_REPLICATION_MASTER_PASSWORD=replication_password,
        )
    )


def _run_subprocess(command, shell=False, env=None, print_stdout=False):
    result = None

    try:
        if shell:
            result = subprocess.check_output(command, shell=True, env=env)
        else:
            result = subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        print(f"\nError running command \"{e.cmd}\", return code \"{e.returncode}\"")
        if e.stdout:
            print("STDOUT:")
            print(e.stdout.decode("utf-8"))
        if e.stderr:
            print("STDERR:")
            print(e.stderr.decode("utf-8"))

    if shell and print_stdout:
        print("COMMAND:")
        print(command)
        print("STDOUT:")
        print(result.decode("utf-8"))

    return result


def _get_num_instances_configured():
    return SLAPD_NUMBER_OF_INSTANCES


def _get_yandex_environment() -> YandexEnvironment:
    try:
        with open(YANDEX_ENVIRONMENT_TYPE_FILE_PATH) as f:
            read_data = f.read().rstrip()
        if read_data == YandexEnvironmentKnownTypes.PRODUCTION:
            return YandexEnvironment.PRODUCTION
        elif read_data == YandexEnvironmentKnownTypes.TESTING:
            return YandexEnvironment.TESTING
        return YandexEnvironment.UNKNOWN
    except FileNotFoundError:
        return YandexEnvironment.UNKNOWN


if __name__ == "__main__":
    args = parse_args()

    if args.cli_action == "get-ldap-replication-password":
        get_replication_password_from_vault(args.yandex_vault_username)
    elif args.cli_action == "setup-ldap-replication":
        setup_ldap_replication(args.replication_master_host)
    elif args.cli_action == "copy-database":
        copy_database(args.database_source_host)
