from argparse import ArgumentParser
from collections import namedtuple

import sys
import pgmigrate
import psycopg2
import yaml

from library.python.vault_client.instances import Production as VaultClient


Config = namedtuple("Config", "host, port, dbname, user, password_secret_id, password_secret_field, options")


def load_database_config(db_config, default_config):
    get_or_default = lambda key : db_config[key] if key in db_config else default_config[key]
    return Config(
        get_or_default("host"),
        get_or_default("port"),
        get_or_default("dbname"),
        get_or_default("user"),
        get_or_default("password_secret_id"),
        get_or_default("password_secret_field"),
        db_config.get("options", None)
    )


def load_config(config_path):
    result = dict()  # database -> Config
    config = yaml.load(open(config_path), Loader=yaml.FullLoader)
    default_database_params = config["default"]
    envs_config = config["environments"]
    for database in envs_config:
        result[database] = load_database_config(envs_config[database], default_database_params)

    return result


def load_yav_token():
    with open('/etc/yandex/maps/wiki/secrets/tokens/robot-wikimap-yav-token', 'r') as f:
        return f.read()


def load_secret(secret_id, field):
    client = VaultClient(
        decode_files=True,
        authorization='OAuth {}'.format(load_yav_token())
    )
    head_version = client.get_version(secret_id)
    return head_version["value"][field]


def call_pgmigrate(db_config, migrations_path, baseline, target_version, verbose, action):
    # call pgmigrate utility from the script
    # emulate sys args for command pgmirate -c {conn_string} -d {base_directory} -t {migration target}
    conn_str = "host={} port={} dbname={} user={} password={} target_session_attrs=read-write".format(
        db_config.host, db_config.port, db_config.dbname, db_config.user,
        load_secret(db_config.password_secret_id, db_config.password_secret_field)
    )
    if (db_config.options):
        conn_str += " " + db_config.options

    if action == 'version':
        conn = psycopg2.connect(conn_str)
        cursor = conn.cursor()
        cursor.execute('SELECT MAX(version) FROM public.schema_version')
        print(cursor.fetchone()[0])
        exit(0)

    sys.argv = [sys.argv[0]] + ['-d', migrations_path, '-c', conn_str]
    if target_version:
        sys.argv += ['-t', target_version]
    if baseline:
        sys.argv += ['-b', baseline]
    if verbose:
        sys.argv.append('-' + 'v' * verbose)
    sys.argv.append(action)
    exit(pgmigrate._main())


def main():
    parser = ArgumentParser(description='run migrations')
    parser.add_argument('-c', '--config', help='path to configuration file', required=True)
    parser.add_argument('-e', '--environment', help='environment', required=True)
    parser.add_argument('-d', '--base_dir', default='.', help='path for migrations')
    parser.add_argument('-b', '--baseline', help='Baseline version')
    parser.add_argument('-t', '--target', help='migration version or "latest"')
    parser.add_argument('-v', '--verbose', default=0, action='count', help='Be verbose')
    parser.add_argument('action', nargs='?', default='info', choices=['migrate', 'info', 'baseline', 'version'])
    argv = parser.parse_args()

    envs_config = load_config(argv.config)
    if argv.environment not in envs_config:
        raise Exception("Environment `{}` is not found in `{}`".format(argv.environment, argv.config))

    db_config = envs_config[argv.environment]

    call_pgmigrate(db_config, argv.base_dir, argv.baseline, argv.target, argv.verbose, argv.action)
