#!/usr/bin/python2

from __future__ import print_function
import subprocess
import argparse
import ctypes
import signal
import errno
import sys
import os
import re

CLOCK_MONOTONIC = 1 # see <linux/time.h>

class timespec(ctypes.Structure):
    _fields_ = [
        ('tv_sec', ctypes.c_long),
        ('tv_nsec', ctypes.c_long)
    ]

librt = ctypes.CDLL('librt.so.1', use_errno=True)
clock_gettime = librt.clock_gettime
clock_gettime.argtypes = [ctypes.c_int, ctypes.POINTER(timespec)]

def monotonic_time():
    t = timespec()
    if clock_gettime(CLOCK_MONOTONIC, ctypes.pointer(t)) != 0:
        errno_ = ctypes.get_errno()
        raise OSError(errno_, os.strerror(errno_))
    return t.tv_sec + t.tv_nsec * 1e-9

class Work(object):
    def __init__(self, prefix, name):
        stem = prefix
        if name:
            stem += '-' + name
        self.data = open(stem + '.time', 'a')
        self.log = open(stem + '.log', 'a')
        self.prepare = None
        self.execute = None
        self.cleanup = None
    def __del__(self):
        print('', file=self.data)
        self.data.close()
        print('', file=self.log)
        self.log.close()
    def clock_start(self):
        self.start_date = monotonic_time()
    def clock_stop(self, rusage):
        end_date = monotonic_time()
        print('real', end_date - self.start_date, file=self.data)
        print('user', rusage.ru_utime, file=self.data)
        print('sys', rusage.ru_stime, file=self.data)

class Monitor(object):
    def __init__(self, prefix, name, execute):
        stem = '%s-%s' % (prefix, name)
        self.logfile = open(stem + '.data', 'a')
        self.execute = execute
    def __del__(self):
        print('', file=self.logfile)
        self.logfile.close()
    def collect(self):
        subprocess.Popen(self.execute, shell=True,
                         stdout=self.logfile)

class Watcher(object):
    def __init__(self, prefix, name, execute):
        stem = '%s-%s' % (prefix, name)
        self.logfile = open(stem + '.data', 'a')
        self.execute = execute
    def __del__(self):
        print('', file=self.logfile)
        self.logfile.close()
    def start(self):
        self.proc = subprocess.Popen(self.execute, shell=True,
                                     stdout=self.logfile)
    def stop(self):
        self.proc.terminate()

defs = {}
works = {}
monitors = []
watchers = []

def collect():
    for m in monitors:
        m.collect()

def sigalrm_handler(sig, frame):
    collect()

def watchers_start():
    for w in watchers:
        w.start()

def watchers_stop():
    for w in watchers:
        w.stop()

parser = argparse.ArgumentParser()
parser.add_argument('-s', '--spec', action='append', default=[])
parser.add_argument('-o', '--oneshot', action='store_true')
parser.add_argument('name')
args = parser.parse_args()

def save_section(type_, name, body):
    if type_ == 'monitor':
        monitors.append(Monitor(args.name, name, body))
    elif type_ == 'watcher':
        watchers.append(Watcher(args.name, name, body))
    else:
        work = works.get(name)
        if not work:
            work = Work(args.name, name)
            works[name] = work
        if type_ == 'prepare':
            work.prepare = body
        elif type_ == 'work':
            work.execute = body
        elif type_ == 'cleanup':
            work.cleanup = body

def resolve_refs(s):
    cre = re.compile('%{([^}]*)}')
    return cre.sub(lambda m: defs[m.group(1)], s)

for s in args.spec:
    type_ = None
    name = None
    body = ''
    for line in open(s):
        if line.startswith('%define'):
            parts = line.split()
            defs[parts[1]] = resolve_refs(parts[2])
        elif line.startswith('%'):
            if type_:
                save_section(type_, name, body)
            parts = line.split()
            type_ = parts[0][1:]
            if len(parts) > 1:
                name = parts[1]
            else:
                name = ''
            body = ''
        elif type_:
            body += resolve_refs(line)
    if type_:
        save_section(type_, name, body)

print("DEBUG: preparing")

noreap = []
for w in works.values():
    if w.prepare:
        sub = subprocess.Popen(w.prepare, shell=True,
                               stdout=w.log, stderr=w.log)
        noreap.append(sub)
while True:
    try:
        os.wait()
    except OSError:
        break

print("DEBUG: working")

collect()

if not args.oneshot:
    signal.signal(signal.SIGALRM, sigalrm_handler)
    signal.setitimer(signal.ITIMER_REAL, 1.0, 1.0)
    watchers_start()

workers = {}
noreap = []
for w in works.values():
    sub = subprocess.Popen(w.execute, shell=True,
                           stdout=w.log, stderr=w.log)
    w.clock_start()
    workers[sub.pid] = w
    noreap.append(sub)
while workers:
    try:
        pid, status, rusage = os.wait3(0)
    except OSError, e:
        if e.errno == errno.ECHILD:
            print('warning: os.wait() missed children!')
            break
        continue
    if not pid:
        continue
    e = os.WEXITSTATUS(status)
    if e:
        print('child exited with status %d' % e)
        sys.exit(1)
    if pid in workers:
        w = workers.pop(pid)
        w.clock_stop(rusage)

if not args.oneshot:
    signal.setitimer(signal.ITIMER_REAL, 0.0, 0.0)
    watchers_stop()

collect()

print('DEBUG: cleaning')

noreap = []
for w in works.values():
    if w.cleanup:
        sub = subprocess.Popen(w.cleanup, shell=True,
                               stdout=w.log, stderr=w.log)
        noreap.append(sub)
while True:
    try:
        os.wait()
    except OSError:
        break