#!/usr/bin/python
#
# Write branches, tellers and accounts file similar to what pgbench creates

import sys
import os
import csv
import argparse
import random
import time

parser = argparse.ArgumentParser()
parser.add_argument("-s", "--scale", help="scale", default=100, type=int)
parser.add_argument("-d", "--dir", help="directory to write data files", default='/tmp')
args = parser.parse_args()

# create sensible dialect
class excel_linefeed(csv.excel):
    lineterminator = '\n'

# write the branches
filename = 'branches.dat'
filepath = os.path.join(args.dir, filename)
fieldnames = [
    'bid',
    'bbalance',
    'filler',
]
fd = open(filepath, 'wb')
stream = csv.DictWriter(fd, fieldnames=fieldnames, dialect=excel_linefeed)
for i in xrange(1, args.scale + 1):
    stream.writerow({
        'bid': i,
        'bbalance': 0,
        'filler': '\N',
    })
fd.close()

# write the tellers
filename = 'tellers.dat'
filepath = os.path.join(args.dir, filename)
fieldnames = [
    'tid',
    'bid',
    'tbalance',
    'filler',
]
fd = open(filepath, 'wb')
stream = csv.DictWriter(fd, fieldnames=fieldnames, dialect=excel_linefeed)
for i in xrange(1, args.scale * 10 + 1):
    j = i / 10 + 1
    if j > args.scale:
        j = args.scale
    stream.writerow({
        'tid': i,
        'bid': j,
        'tbalance': 0,
        'filler': '\N',
    })
fd.close()

# write the accounts
filename = 'accounts.dat'
filepath = os.path.join(args.dir, filename)
fieldnames = [
    'aid',
    'bid',
    'abalance',
    'filler',
]
fd = open(filepath, 'wb')
stream = csv.DictWriter(fd, fieldnames=fieldnames, dialect=excel_linefeed)
for i in xrange(1, args.scale * 100000 + 1):
    j = i / 100000 + 1
    if j > args.scale:
        j = args.scale
    stream.writerow({
        'aid': i,
        'bid': j,
        'abalance': 0,
        'filler': 84*' ',
    })
fd.close()