#! /usr/bin/env python3
import subprocess, time, os, sys, argparse, logging

log = logging.getLogger('set-admin-password')


def get_options():
    parser = argparse.ArgumentParser(
        description='Ensure a user is available and set password'
    )
    parser.add_argument('--user', type=str, help='Database username')
    parser.add_argument('--password', type=str, help='Database password')
    parser.add_argument('--name', type=str, help='Database name')
    parser.add_argument(
        '--stats-name', type=str, help='Stats-database name', default=None
    )
    return parser


def su_pg(command, db='postgres') -> str:
    if isinstance(command, str):
        command = command.encode('utf-8')
    return subprocess.check_output(
        f'su postgres -c \'psql -P pager=off  -t -A {db}\'',
        input=command,
        shell=True,
        env=os.environ.copy(),
    ).decode('utf-8')


def main():
    options = get_options().parse_args()
    if not options.user:
        raise SystemExit(1)
    for i in range(30):
        if subprocess.call(['pg_isready']) == 0:
            break
        else:
            log.info("Waiting for postgresql to be ready")
            time.sleep(1.0)
    users = [x for x in su_pg("select usename from pg_user").splitlines()]
    if options.user not in users:
        log.info("Creating user %s", options.user)
        su_pg(f'CREATE USER {options.user} WITH LOGIN PASSWORD \'{options.password}\'')
    else:
        su_pg(f'ALTER USER {options.user} WITH LOGIN PASSWORD \'{options.password}\'')
    dbs = [x for x in su_pg("select datname from pg_database").splitlines()]
    time.sleep(1)
    if options.name and options.name not in dbs:
        log.info("Creating db %s", options.name)
        su_pg(
            f'CREATE DATABASE {options.name} WITH ENCODING \'utf-8\' OWNER {options.user}'
        )
    if options.stats_name and options.stats_name not in dbs:
        log.info("Creating stats db %s", options.stats_name)
        su_pg(
            f'CREATE DATABASE {options.stats_name} WITH ENCODING \'utf-8\' OWNER {options.user}'
        )
        su_pg(
            "CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE",
            options.stats_name,
        )


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    main()
