#!/usr/bin/env python3

import re
import os
import subprocess
import statistics


class NoQueryPlanOutput(Exception):
    pass


def main(dbname, sql_file, times=10):
    if not os.path.isfile(sql_file):
        raise FileNotFoundError(sql_file)

    with open(sql_file) as f:
        sql = f.read()
    sql = '\n'.join(
        x for x in sql.splitlines()
        if not x.strip().startswith('--')
    )
    if not sql.lower().strip().startswith('explain '):
        # In case you forgot to put the 'EXPLAIN ANALYZE' before the query
        sql = 'EXPLAIN ANALYZE {}'.format(sql)

    execution_times = []
    planning_times = []
    for i in range(times):
        p = subprocess.Popen(
            ['psql', dbname],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out, err = p.communicate(sql.encode('utf-8'))
        if err:
            raise Exception(err.decode('utf-8'))
        output = out.decode('utf-8')
        if 'QUERY PLAN' not in output:
            raise NoQueryPlanOutput(output)

        planning_time, = re.findall('Planning time: ([\d\.]+) ms', output)
        execution_time, = re.findall('Execution time: ([\d\.]+) ms', output)
        planning_times.append(float(planning_time))
        execution_times.append(float(execution_time))

    def fmt(ms):
        return '{:.2f}ms'.format(ms)

    print("Execution time")
    print(
        "\tBEST   ",
        fmt(min(execution_times)).ljust(15),
        '<-- most important number'
    )
    print("\tmean   ", fmt(statistics.mean(execution_times)))
    print("\tmedian ", fmt(statistics.median(execution_times)))
    print("Planning time")
    print("\tbest   ", fmt(min(planning_times)))
    print("\tmean   ", fmt(statistics.mean(planning_times)))
    print("\tmedian ", fmt(statistics.median(planning_times)))


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Best EXPLAIN ANALYZE')
    parser.add_argument('dbname', type=str,
                        help='Database name (e.g. `psql $this`)')
    parser.add_argument('sql_file', type=str,
                        help='The .sql file (e.g. `psql $dbname < $this`)')
    parser.add_argument('--times', dest='times', type=int,
                        default=10,
                        help='Number of times to run (default 10)')

    args = parser.parse_args()
    args = vars(args)
    main(args.pop('dbname'), args.pop('sql_file'), **args)