#!/usr/bin/python

import os
import sys
# we use gdbm instead of anydbm because IME, berkeley db gets corrupted if a filesystem fills up and anydbm appears to default to berkeley db
import gdbm

try:
    import gmpy
except:
    use_gmpy = False
else:
    use_gmpy = True

try:
    import cPickle
except:
    import pickle
else:
    pickle = cPickle

import rangeset

sys.path.insert(0, '/usr/local/lib')
sys.path.insert(0, os.path.expanduser('~/lib'))

initialize = False
run = False

database = None

hosts = []

def cvt(x):
    if use_gmpy:
        return gmpy.mpz(x)
    else:
        return int(x)

offset = cvt(0)

def usage(retval):
    sys.stderr.write('Usage: %s\n' % sys.argv[0])
    sys.stderr.write('\t--database dbname\n')
    sys.stderr.write('\t--initialize command offset low high host1 host2 ... hostn\n')
    sys.stderr.write('\t--run\n')
    sys.stderr.write('\t--help\n')
    sys.exit(retval)

while sys.argv[1:]:
    if sys.argv[1] == '--database':
        database = sys.argv[1]
    elif sys.argv[1] == '--initialize':
        initialize = True
        command = sys.argv[2]
        offset = sys.argv[3]
        low = cvt(sys.argv[4])
        high = cvt(sys.argv[5])
        hosts = sys.argv[6:]
        break
    elif sys.argv[1] == '--run':
        run = True
    elif sys.argv[1] in [ '-h', '--help' ]:
        usage(0)
    else:
        sys.stderr.write('%s: Unrecognized option: %s\n' % (sys.argv[0], sys.argv[1]))
        usage(1)
    del sys.argv[1]

if database == None:
    sys.stderr.write('%s: --database required before --initialize\n' % sys.argv[1])
    usage(1)

if initialize:
    desired = open('%s-desired.pickle', 'w')
    rs = rangeset.rangeset()
    rs.add(low - offset, high - offset)
    pickle.dump(rs, desired)
    desired.close()

    db = gdbm.open('%s-results.gdbm' % database, 'w')
    db.close()

    meta = open('%s-meta.txt' % database, 'w')
    meta.write('command %s\n' % command)
    meta.write('offset %s\n' % str(offset))
    meta.write('hosts %s\n' % ','.join(hosts))
    meta.close()

def read_meta(meta, kind, numeric=False, separator=None):
    line = meta.readline()
    fields = line.split()
    if fields[0] != kind:
        sys.stderr.write('%s: Unexpected line in meta file: wanted %s, got %s\n' % (sys.argv[0], kind, fields[0]))
        sys.exit(1)
    if numeric:
        return cvt(line.partition(' ')[2])
    elif separator != None:
        return ' '.join(fields[1:])
    else:
        return line.partition(' ')[2]
    
if run:
    desired = open('%s-desired.pickle', 'r+')
    pickle.load(rs, desired)
    desired.close()

    # Assume we're restarting - nothing previously invoked is still running.
    # If there's anything in the in-progress file, that means a previous computation was interrupted - move these numbers
    # back to desired for now
    in_progress = gdbm.open('%s-in-progress.gdbm', 'w+')
    changed = False
    for key in in_progress.keys():
        changed = True
        rs.add(key, key)

    in_progress.close()

    if changed:
        desired = open('%s-desired.pickle', 'w')
        pickle.dump(rs, desired)
        desired.close()

    in_progress = gdbm.open('%s-in-progress.gdbm', 'c')

    meta = open('%s-meta.txt' % database, 'r')
    command = read_meta(meta, 'command')
    # may as well do this one as a string - we'll just need it as a string again soon anyway
    offset = read_meta(meta, 'offset')
    hosts = read_meta(meta, 'hosts', separator=',')
    meta.close()

    busyness = {}