#!/usr/bin/python3

import socket
import os
import tempfile
import shutil
import subprocess
import pytest
import sys
import signal
import time
import psycopg2

from os.path import join, exists

lib_postgresql = '/usr/lib/postgresql/'
host = 'localhost'

def get_unused_port():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind((host, 0))
    port = sock.getsockname()[1]
    sock.close()

    return port

def find_postgresql_bin():
    versions = [int(d) for d in os.listdir(lib_postgresql) if d.isdigit()]

    for v in sorted(versions, reverse=True):
        bin_dir = join(lib_postgresql, str(v))
        if all(exists(join(bin_dir, 'bin', f)) for f in ('initdb', 'postgres')):
            return bin_dir

def clean_up(base_dir):
    print('clean up')
    shutil.rmtree(base_dir, ignore_errors=True)

def run_initdb(bin_dir, base_dir):
    print('running initdb')
    initdb = join(bin_dir, 'bin', 'initdb')

    args = [initdb,
            '-D', join(base_dir, 'data'),
            '--lc-messages=C',
            '-U', 'postgres',
            '-A', 'trust']
    subprocess.run(args, check=True)

def run_server(bin_dir, base_dir, port):
    print('starting server')
    postgres = join(bin_dir, 'bin', 'postgres')
    assert os.path.exists(postgres)

    args = [postgres,
            '-p', str(port),
            '-D', os.path.join(base_dir, 'data'),
            '-k', os.path.join(base_dir, 'tmp'),
            '-h', host,
            '-F',
            '-c', 'logging_collector=off']
    return subprocess.Popen(args, stderr=subprocess.PIPE)

def stop_server(server):
    print('stopping server')
    if server is None:
        return

    server.send_signal(signal.SIGTERM)
    t0 = time.time()

    while server.poll() is None:
        if time.time() - t0 > 5:
            server.kill()
        time.sleep(0.1)

def configure_database(port):
    conn_params = {'user': 'postgres', 'host': host, 'port': port}
    for _ in range(50):
        try:
            conn = psycopg2.connect(dbname='template1', **conn_params)
        except psycopg2.OperationalError as e:
            allowed = [
                'the database system is starting up',
                'could not connect to server: Connection refused',
                'failed: Connection refused',
            ]
            if not any(msg in e.args[0] for msg in allowed):
                raise
            time.sleep(0.1)
            continue
        break
    conn.set_session(autocommit=True)
    cur = conn.cursor()
    cur.execute("CREATE ROLE gis PASSWORD 'gis' SUPERUSER CREATEDB CREATEROLE INHERIT LOGIN")
    cur.execute("CREATE DATABASE gis")
    cur.execute('GRANT CREATE ON DATABASE gis TO "gis"')
    conn.close()

    conn = psycopg2.connect(dbname='gis', **conn_params)
    conn.set_session(autocommit=True)
    cur = conn.cursor()
    cur.execute('CREATE SCHEMA gis')
    cur.execute('GRANT USAGE,CREATE ON SCHEMA gis TO "gis"')
    cur.execute('CREATE EXTENSION postgis')
    cur.execute('CREATE EXTENSION postgis_raster')
    cur.execute('SELECT postgis_version()')
    cur.fetchone()
    conn.close()

def main():
    bin_dir = find_postgresql_bin()

    if not bin_dir:
        print('postgresql bin dir not found')
        sys.exit(1)

    port = get_unused_port()

    base_dir = tempfile.mkdtemp()
    os.mkdir(join(base_dir, 'tmp'))

    run_initdb(bin_dir, base_dir)

    server = run_server(bin_dir, base_dir, port)

    configure_database(port)

    os.environ['PGPORT'] = str(port)
    ret = pytest.main()
    stop_server(server)

    server_errors = server.stderr.read()
    if ret:
        print('server error log')
        print()
        print(server_errors.decode('utf-8'))

    clean_up(base_dir)


if __name__ == '__main__':
    if os.getuid() == 0:  # Postgres refuses to run as root
        sys.exit(0)

    main()
    sys.exit(0)
