#!/bin/sh
"""": # -*-python-*-
command -v python3 > /dev/null && exec python3 "$0" "$@"
command -v python2 > /dev/null && exec python2 "$0" "$@"
echo "error: unable to find python3 or python2" 1>&2; exit 2
"""
from __future__ import print_function

import sys
py3 = sys.version_info.major > 2

from argparse import ArgumentParser, RawDescriptionHelpFormatter
from base64 import b64encode
from collections import namedtuple
from errno import ENOENT
from os import chdir, environ, fdopen, getcwd, mkdir, symlink, umask, urandom
from os.path import abspath, basename, isdir, isfile, islink
from shutil import rmtree
from sys import stderr, stdout
from subprocess import PIPE, Popen
import argparse, os, re, sys
if py3:
    from shlex import quote
else:
    from pipes import quote

version = (0, 0, 0)
version_str = '.'.join(str(x) for x in version)

if py3:
    def iteritems(x):
        return x.items()
else:
    def iteritems(x):
        return x.iteritems()

ex_res = namedtuple('SubprocResult', ['out', 'err', 'proc', 'rc'])

# The run, logcmd, ex, and exo commands were copied from bup by the author.

def run(cmd, check=True, input=None, **kwargs):
    """Run a subprocess as per subprocess.Popen(cmd, **kwargs) followed by
    communicate(input=input).  If check is true, then throw an
    exception if the subprocess exits with non-zero status.  Return a
    SubprocResult tuple.

    """
    if input:
        assert 'stdin' not in kwargs
        kwargs['stdin'] = PIPE
    p = Popen(cmd, **kwargs)
    out, err = p.communicate(input=input)
    if check and p.returncode != 0:
        raise Exception('subprocess %r failed with status %d%s'
                        % (' '.join(map(quote, cmd)), p.returncode,
                           (', stderr: %r' % err) if err else ''))
    return ex_res(out=out, err=err, proc=p, rc=p.returncode)

def logcmd(cmd):
    if isinstance(cmd, basestring):
        print(cmd, file=stderr)
    else:
        print(' '.join(map(quote, cmd)), file=stderr)

def ex(cmd, **kwargs):
    """Print cmd to stderr and then run it as per ex(...).
    Print the subprocess stderr to stderr if stderr=PIPE and there's
    any data.
    """
    if verbosity > 0:
        logcmd(cmd)
    result = run(cmd, **kwargs)
    if result.err:
        stderr.write(result.err)
    return result

def exo(cmd, **kwargs):
    """Print cmd to stderr and then run it as per ex(..., stdout=PIPE).
    Print the subprocess stderr to stderr if stderr=PIPE and there's
    any data.

    """
    assert 'stdout' not in kwargs
    kwargs['stdout'] = PIPE
    return ex(cmd, **kwargs)

def open_new_private(file, mode):
    assert('w' in mode)
    return fdopen(os.open(file,
                          os.O_WRONLY | os.O_CREAT | os.O_EXCL,
                          0o0600),
                  mode)

_initdb_ver_rx = re.compile(r'^initdb \(PostgreSQL\) (\d+)(?:\.(\d+))?(?:\.(\d+))?(?:$| )'.encode('ascii'))
def initdb_ver(path):
    res = exo((path, '--version'), check=False)
    if res.rc != 0:
        return None
    match = _initdb_ver_rx.match(res.out)
    if match:
        return tuple(int(x) for x in match.groups() if x is not None)


def setup_pg_data(pgbin, pgver, dir, passfile, network, port, bind_addrs, initdb_args):
    ex((pgbin + '/initdb',
        '--pgdata', dir,
        '--username', 'postgres',
        '--auth', 'password',
        '--pwfile', passfile) + tuple(initdb_args))
    mkdir(dir + '/sock')
    # postgresql.conf
    with open(dir + '/postgresql.conf', 'a') as f:
        if pgver[:2] == (8, 4):
            print("unix_socket_directory = 'sock'", file=f)
        else:
            print("unix_socket_directories = 'sock'", file=f)
        print("external_pid_file = 'sandbox.pid'", file=f)
        print("listen_addresses = '%s'" % (','.join(bind_addrs)),
              file=f)
        print('port =', port, file=f)

def require_pgbin_val(val):
    if not val:
        print('PostgreSQL bin directory not specified', file=stderr)
        sys.exit(1)

def require_sandbox_val(val):
    if not val:
        print('sandbox directory not specified', file=stderr)
        sys.exit(1)

def require_sandbox(dir):
    if not (isdir(dir + '/data')
            and isfile(dir + '/pgpass')
            and islink(dir + '/pg-bin')):
        print('%r does not appear to be a PostgreSQL sandbox' % dir,
              file=stderr)
        sys.exit(1)

def cmd_init(args):
    require_pgbin_val(args.pgbin)
    require_sandbox_val(args.sandbox)
    pgbin = abspath(args.pgbin)
    pgver = initdb_ver(pgbin + '/initdb')
    sandbox = abspath(args.sandbox)
    admin_pw = b64encode(urandom(32)).decode('ascii')
    mkdir(sandbox)
    try:
        with open_new_private(sandbox + '/pass-admin', 'w') as f:
            print(admin_pw, file=f, end='')
        with open_new_private(sandbox + '/pgpass', 'w') as f:
            print('# hostname:port:database:username:password', file=f)
            print('*:*:*:*:' + admin_pw, file=f)
        if args.network:
            with open(sandbox + '/port', 'w') as f:
                print(args.port, file=f)
            with open(sandbox + '/bind-addrs', 'w') as f:
                for addr in args.bind_addr:
                    print(addr, file=f)
        initdb_args = args.initdb_args
        if initdb_args and initdb_args[0] == '--':
            initdb_args = initdb_args[1:]
        setup_pg_data(pgbin, pgver, sandbox + '/data', sandbox + '/pass-admin',
                      args.network, args.port, args.bind_addr, initdb_args)
        symlink(pgbin, sandbox + '/pg-bin')
        mkdir(sandbox + '/ext')
        with open(sandbox + '/pg-sandbox', 'w') as f:
            print('created by pgbox', version_str, file=f)
    except (Exception, KeyboardInterrupt, SystemExit) as ex:
        rmtree(sandbox)
        raise

def get_port(sandbox):
    port = None
    try:
        with open(sandbox + '/port', 'r') as f:
            port = int(f.read())
    except IOError as ex:
        if ex.errno != ENOENT:
            raise
    return port

def get_host(sandbox):
    host = None
    try:
        with open(sandbox + '/bind-addrs', 'r') as f:
            addrs = f.readlines()
            host = addrs[-1].strip()
    except IOError as ex:
        if ex.errno != ENOENT:
            raise
    return host

def configure_env(sandbox, overrides):
    port = get_port(sandbox)
    host = sandbox + '/data/sock'
    if port:
        host = get_host(sandbox)
        if not host:  # backwards compatibility
            host = 'localhost'
            with open(sandbox + '/bind-addrs', 'w') as f:
                print(host, file=f)
    os.environ['PGBOX'] = sandbox
    os.environ['PGPASSFILE'] = sandbox + '/pgpass'
    os.environ['PGHOST'] = host
    os.environ['PGPORT'] = str(port)
    os.environ['PATH'] = sandbox + '/pg-bin:' + os.environ['PATH']
    os.environ['PGDATA'] = sandbox + '/data'
    for name, val in iteritems(overrides):
        os.environ[name] = val

def cmd_env(args):
    require_sandbox_val(args.sandbox)
    sandbox = abspath(args.sandbox)
    require_sandbox(sandbox)
    env_args = []
    env_overrides = {}
    for arg in args.env_args:
        if '=' not in arg:
            env_args = args.env_args[len(env_overrides):]
            break
        name, val = arg.split('=', 1)
        env_overrides[name] = val
    configure_env(sandbox, env_overrides)
    if not args.env_cmd:
        for item in iteritems(os.environ):
            print('%s=%s' % item)
        return
    stdout.flush()
    stderr.flush()
    exe = args.env_cmd
    os.execlp(exe, basename(exe), *env_args)

def cmd_version(args):
    print('pgbox ' + version_str)


# Because earlier versions of argpase will report "error: too few
# arguments" when there's no subcommand, but we want "pgbox --version"
# to work.
if sys.argv[1:] == ['--version']:
    cmd_version([])
    exit(0)

description = """
Create PostgreSQL sandboxes and then run commands that use them.
"""

epilog = """
EXAMPLES:
  pgbox init --pgbin /usr/lib/postgresql/10/bin \\
    --port=12345 --box ./box --  -E UTF8 --locale=C
  export PGBOX="$(pwd)/box"
  pgbox env pg_ctl start -o -F
  pgbox env pg_ctl stop

"""

# https://docs.python.org/2/library/argparse.html#sub-commands
parser = ArgumentParser(prog='pgbox', description=description, epilog=epilog,
                        formatter_class=RawDescriptionHelpFormatter)
parser.add_argument('--version', action='store_true', help='print the %s version' % parser.prog)
parser.add_argument('--verbose', '-v', action='count', default=0)
sp = parser.add_subparsers()

# pgbox version
ver_sp = sp.add_parser('version', help='print the %s version' % parser.prog)
ver_sp.set_defaults(action=cmd_version)

# pgbox init
init_sp = sp.add_parser('init', help='initialize a sandbox')
init_sp.add_argument('--port', type=int, default=5432)
init_sp.add_argument('--network', action='store_true', default=True)
init_sp.add_argument('--no-network', action='store_false', dest='network')
init_sp.add_argument('--pgbin', metavar='DIRECTORY', required=False,
                     default=environ.get('PGBOX_PGBIN'))
init_sp.add_argument('--sandbox', '-s', required=False,
                     default=environ.get('PGBOX'))
init_sp.add_argument('--bind-addr', '-b', action='append', dest='bind_addr')
init_sp.add_argument('initdb_args', nargs=argparse.REMAINDER,
                     metavar='[-- INITDB_ARG...]')
init_sp.set_defaults(action=cmd_init)

# pgbox env
init_sp = sp.add_parser('env', help='run command in a sandbox environment')
init_sp.add_argument('--sandbox', '-s', required=False,
                     default=environ.get('PGBOX'))
init_sp.add_argument('env_cmd', nargs='?', metavar='COMMAND')
#init_sp.add_argument('env_args', nargs='*', metavar='ARG')
init_sp.add_argument('env_args', metavar='ARG...', nargs=argparse.REMAINDER)
init_sp.set_defaults(action=cmd_env)

args = parser.parse_args()
verbosity = args.verbose
if not getattr(args, 'bind_addr', None):
    args.bind_addr = ['localhost']

if args.version:
    if hasattr(args, 'action') and args.action != cmd_version:
        parser.print_usage()
        exit(2)
    cmd_version([])
    exit(0)

if not hasattr(args, 'action'):
    parser.print_usage()
    exit(2)

args.action(args)