#!/usr/bin/env python3
"""citus_dev

Usage:
  citus_dev make <name> [--size=<count>] [--port=<port>] [--use-ssl] [--no-extension] [--destroy] [--init-with=<sql_file>] [--init-worker-with=<sql_file>] [--with-pgbouncer]
  citus_dev restart <name> [--watch] [--port=<port>]
  citus_dev (start|stop) <name> [--port=<port>] [--force]

Options:
  --size=<count>           Number of workers to create when 0 the coordinator will be added as a worker [default: 2]
  --port=<port>            Port number to use for the coordinator. All workers take subsequent numbers [default: 9700]
  --watch                  Watch for changes to the citus plugin and restart the cluster when the plugin updates
  --use-ssl                Create the cluster with ssl enabled
  --no-extension           Do not create the extension while creating the nodes
  --destroy                Destroy any old cluster with the same name
  --init-with=<sql_file>   A SQL script to run after creation of the cluster to set up any necessary tables and data
  --init-worker-with=<sql_file>   A SQL script to run after creation of the cluster on the workers
  --with-pgbouncer         Setup pgbouncers between worker and coordinator (requires citus enterprise)
  --force                  Forceful shutdown

"""
from docopt import docopt
from subprocess import call
from subprocess import Popen, PIPE
import os
import subprocess
import sys
import getpass
import time
import distutils.spawn

# for osx we might want to start postgres via a fixopen binary that preloads a
# dylib to fix the interrupted systemcall, this is done by changing the postgres
# path for pg_ctl to the fixopen binary
pgctl_flags = ""
fixopen = distutils.spawn.find_executable("postgres.fixopen")
if fixopen:
    pgctl_flags += f' -p "{fixopen}"'

def run(command, *args, **kwargs):
    print(command)
    result = subprocess.run(command, *args, check=True, shell=True, **kwargs)
    print()
    return result


def createNodeCommands(clustername, role, index=None, usessl=False):
    nodename = role
    if index != None:
        nodename += "%d" % index

    run(f"initdb -D {clustername}/{nodename}")

    shared_preload_libraries = ['citus', 'pg_stat_statements']

    shared_preload_libraries = ','.join(shared_preload_libraries)

    run(f"echo \"shared_preload_libraries = '{shared_preload_libraries}'\" >> {clustername}/{nodename}/postgresql.conf")
    run(f'echo "wal_level = logical" >> {clustername}/{nodename}/postgresql.conf')

    if usessl:
        run(f'echo "ssl = on" >> {clustername}/{nodename}/postgresql.conf')
        run(f"echo \"citus.node_conninfo = 'sslmode=require'\" >> {clustername}/{nodename}/postgresql.conf")
        run(f"openssl req -new -x509 -days 365 -nodes -text -out {clustername}/{nodename}/server.crt -keyout {clustername}/{nodename}/server.key -subj '/CN={nodename}'")
        run(f"chmod 0600 {clustername}/{nodename}/server.key")

def createPgBouncerUsers(clustername):
    username = getpass.getuser()
    with open(f"{clustername}/users.txt", "w") as f:
        f.write(f'"{username}" ""')

def createPgBouncerConfig(clustername, port, index):
    workerPort = port + index + 1
    bouncerPort = port + index + 101
    username = getpass.getuser()

    bouncerConfig = f"""
[databases]
postgres = host=127.0.0.1 port={workerPort}

[pgbouncer]
application_name_add_host = 1
pool_mode = transaction
listen_port = {bouncerPort}
listen_addr = *
auth_type = trust
auth_file = {clustername}/users.txt
logfile = worker{index}.pgbouncer.log
pidfile = {clustername}/worker{index}.pgbouncer.pid
admin_users = {username}
stats_users = {username}
client_tls_key_file = {clustername}/worker{index}/server.key
client_tls_cert_file = {clustername}/worker{index}/server.crt
client_tls_sslmode = prefer
"""

    with open(f"{clustername}/worker{index}.pgbouncer.ini", "w") as f:
        f.write(bouncerConfig)


def main(arguments):
    print(arguments)
    if arguments["make"]:
        clustername = arguments["<name>"]

        if arguments['--destroy']:
            stopCluster(clustername, True)
            run(f'rm -rf {clustername}')


        createNodeCommands(
            clustername,
            "coordinator",
            usessl=arguments["--use-ssl"],
        )

        port = int(arguments["--port"])
        size = int(arguments["--size"])
        pgbouncer = bool(arguments["--with-pgbouncer"])

        if pgbouncer:
            createPgBouncerUsers(clustername)

        for i in range(size):
            createNodeCommands(
                arguments["<name>"],
                "worker",
                i,
                usessl=arguments["--use-ssl"],
            )
            if pgbouncer:
                createPgBouncerConfig(clustername, port, i)

        cport = port
        role = "coordinator"
        run(f'pg_ctl {pgctl_flags} -D {clustername}/{role} -o "-p {cport}" -l {role}_logfile start')
        port += 1

        worker_ports = []
        for i in range(size):
            role = "worker%d" % i
            worker_ports.append(port)
            run(f'pg_ctl {pgctl_flags} start -D {clustername}/{role} -o "-p {port}" -l {role}_logfile')
            port += 1
        port = cport

        if getpass.getuser() != 'postgres' and not os.getenv('PGDATABASE'):
            for i in range(size + 1):
                nodeport = port + i
                run(f'createdb -p {nodeport}')

        if not arguments["--no-extension"]:
            for i in range(size + 1):
                nodeport = port + i
                run(f'psql -p {nodeport} -c "CREATE EXTENSION citus;"')

            run(f"psql -p {port} -c \"SELECT * from citus_set_coordinator_host('localhost', {port});\"")

            for i in range(size):
                workerport = port + 1 + i
                run(f"psql -p {port} -c \"SELECT * from master_add_node('localhost', {workerport});\"")
            run(f'psql -p {port} -c "SELECT * from master_get_active_worker_nodes();"')

        if pgbouncer:
            # need to start pgbouncers and configure pg_dist_poolinfo
            for i in range(size):
                coordinatorPort = port
                workerPort = port + i + 1 
                bouncerPort = port + i + 101
                run(f'pgbouncer -d {clustername}/worker{i}.pgbouncer.ini')
                run(f"psql -p {coordinatorPort} -c \"INSERT INTO pg_dist_poolinfo SELECT nodeid, 'host=localhost port={bouncerPort}' AS poolinfo FROM pg_dist_node WHERE nodeport = {workerPort};\"")


        if arguments['--init-with']:
            run(f'psql -p {cport} -f {arguments["--init-with"]} -v ON_ERROR_STOP=1')
        if arguments['--init-worker-with']:
            for i in range(size):
                workerport = port + 1 + i
                run(f'psql -p {workerport} -f {arguments["--init-worker-with"]} -v ON_ERROR_STOP=1')

    elif arguments["stop"]:
        clusterName = arguments["<name>"]
        always = arguments["--force"]
        stopCluster(clusterName, always)


    elif arguments["start"]:
        clustername = arguments["<name>"]
        port = int(arguments["--port"])
        for role in getRoles(clustername):
            run(f'pg_ctl {pgctl_flags} start -D {clustername}/{role} -o "-p {port}" -l {role}_logfile')
            port += 1
        for bouncerConfig in getPgBouncerConfigs(clustername):
            run(f'pgbouncer -d {clustername}/{bouncerConfig}')



    elif arguments["restart"]:
        clustername = arguments["<name>"]
        port = int(arguments["--port"])
        if arguments["--watch"]:
            run(
                "fswatch -0 '%s' | xargs -0 -n 1 -I{} citus_dev restart %s --port=%d"
                % (citus_so(), clustername, port)
            )

        else:
            cport = port
            for role in getRoles(clustername):
                run(f'pg_ctl {pgctl_flags} restart -D {clustername}/{role} -o "-p {cport}" -l {role}_logfile')
                cport += 1


    else:
        print("unknown command")
        exit(1)

def getPgBouncerPort(clustername, configfile):
    with open(f"{clustername}/{configfile}") as f:
            for line in f.readlines():
                if line.find("listen_port") >= 0:
                    lineParts = line.split("=")
                    return int(lineParts[1])

def stopCluster(clustername, always=False):
    for bouncerConfig in getPgBouncerConfigs(clustername):
        bouncerPort = getPgBouncerPort(clustername, bouncerConfig)
        # the way we shut down pgbouncer always results in an error, so we pipe to true
        run(f"psql -p {bouncerPort} pgbouncer -c \"SHUTDOWN;\" || true")

    pipeTrue = ""
    if always:
        pipeTrue = " || true"

    for role in getRoles(clustername):
        run(f"pg_ctl {pgctl_flags} stop -D {clustername}/{role} {pipeTrue}")

def getPgBouncerConfigs(clustername):
    try:
        bouncerFiles = [f.name for f in os.scandir(clustername) if f.name.find("pgbouncer.ini") >= 0]
        return bouncerFiles
    except FileNotFoundError:
        return []


def getRoles(name):
    try:
        roles = [f.name for f in os.scandir(name) if f.is_dir()]
        roles.sort()
        return roles
    except FileNotFoundError:
        return []


def pg_libdir():
    process = Popen(["pg_config"], stdout=PIPE)
    (output, err) = process.communicate()
    exit_code = process.wait()

    output = str(output)

    for line in output.split("\\n"):
        if line.startswith("LIBDIR"):
            return line.split("=", 1)[1].strip()

    raise Exception("can't find postgres lib dir")


def citus_so():
    return pg_libdir() + "/citus.so"


if __name__ == "__main__":
    print(sys.argv)
    main(docopt(__doc__, version="citus_dev"))
