#!/usr/bin/env python3

import csv
import psycopg2
import logging
import argparse

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

parser = argparse.ArgumentParser()
parser.add_argument("--output", "-o")
parser.add_argument("--iters", "-n", default=10, type=int)
parser.add_argument("--points", "-p", default=10000, type=int)
parser.add_argument("--port", default=5432, type=int)
parser.add_argument("inputs", nargs="*")
args = parser.parse_args()

conn = psycopg2.connect(f'host=/tmp user=bernd dbname=bernd port={args.port}')
cur = conn.cursor()

output = open(args.output, 'w')
out = csv.writer(output)
out.writerow(['schema', 'tbl', 'i', 'point', 'time', 'hit', 'read', 'iotime'])

for schema, tbl in [i.split(".") for i in args.inputs]:
    logging.info(f"Fetching points for {schema}.{tbl}")

    cur.execute(f"SELECT aid, abalance FROM {schema}.{tbl} ORDER BY RANDOM() LIMIT {args.points}")
    datapoints = cur.fetchall()

    logging.info(f"Benching {schema}.{tbl}")
    for iteration in range(args.iters):
        logging.info(f" iter {iteration}")
        for pointid, point in enumerate(datapoints):
            cur.execute(f"EXPLAIN (BUFFERS, ANALYZE, FORMAT 'json') SELECT * FROM {schema}.{tbl} WHERE aid = %s AND abalance = %s",  point)
            explain = cur.fetchone()[0][0]
            plan = explain['Plan']
            exec_time = explain['Execution Time']
            hit = plan['Shared Hit Blocks']
            read = plan['Shared Read Blocks']
            readtime = plan['I/O Read Time']
            out.writerow(list(map(str, [schema, tbl, iteration+1, pointid+1, exec_time, hit, read, readtime])))