#!/usr/bin/env python

import random
import psycopg2
import time
import psutil
import sys
import subprocess

conn = None

def get_mem_usage():

	result = subprocess.run(['./get-mem.sh'], stdout=subprocess.PIPE)
	return result.stdout.decode('utf-8').strip()


def get_timing():
	result = subprocess.run(['./get-timing.sh'], stdout=subprocess.PIPE)
	return result.stdout.decode('utf-8').strip()

def generate_schema(seed, ntables, ncols, indexes = False):

	conn = psycopg2.connect('host=localhost user=tomas dbname=test')

	cur = conn.cursor()

	for t in range(0, ntables):
		cols = [f'c_{t}_{c} int' for c in range(0, ncols)]
		cols = ', '.join(cols)
		cur.execute(f'DROP TABLE IF EXISTS t_{t};')
		cur.execute(f'CREATE TABLE t_{t} ({cols});')

		if indexes:
			for c in range(0, ncols):
				cur.execute(f'CREATE INDEX ON t_{t} (c_{t}_{c});')

	cur.execute('commit')

def generate_data(seed, ntables, ncols):

	conn = psycopg2.connect('host=localhost user=tomas dbname=test')

	cur = conn.cursor()

	for t in range(0, ntables):
		cols = ', '.join(['i/100' for c in range(0, ncols)])
		cur.execute(f'INSERT INTO  t_{t} SELECT {cols} FROM generate_series(1,10000) s(i);')

	cur.execute('analyze')
	cur.execute('commit')


def generate_query(seed, ntables, ncols, nclauses, indexes, data):

	global peak_mem_usage

	conn = psycopg2.connect('host=localhost user=tomas dbname=test')

	configure_session(conn)

	random.seed(seed)

	sql = 'EXPLAIN SELECT * FROM t_0'

	for t1 in range(1, ntables):

		n = random.randint(1,nclauses)
		clauses = []

		for c in range(0, n):

			# random preceding table
			t2 = random.randint(0, t1 - 1)

			# random columns
			c1 = random.randint(0, ncols-1)
			c2 = random.randint(0, ncols-1)

			clauses.append(f'(t_{t1}.c_{t1}_{c1} = t_{t2}.c_{t2}_{c2})')

		clauses = ' AND '.join(clauses)

		join_type = random.choice(['JOIN', 'LEFT JOIN', 'RIGHT JOIN', 'FULL JOIN'])

		sql += f'\n    {join_type} t_{t1} ON ({clauses})'

	cur = conn.cursor()

	try:
		s = time.time()
		cur.execute(sql)
		e = time.time()

		mem = get_mem_usage()
		timing = get_timing()

		print("RESULT", seed, indexes, data, ntables, ncols, nclauses, mem, timing, 1000 * (e - s))

	except Exception as ex:

		print("CRASH", seed, indexes, data, ntables, ncols, nclauses)
		print(ex)
		print(sql)
		time.sleep(5)

	sys.stdout.flush()


def configure_session(conn):

	cur = conn.cursor()
	cur.execute(f'set join_collapse_limit = 32')
	cur.execute(f'set geqo = off')


if __name__ == '__main__':

	for indexes in [True, False]:
		for data in [True, False]:

			generate_schema(0, 16, 32, indexes)

			if data:
				generate_data(0, 16, 32)

			seed = 0
			for ntables in range(2, 32):
				for s in range(0, 1000):
					seed += 1
					generate_query(seed, ntables, 32, 3, data, indexes)

