#!/usr/bin/python

import cliapp
import os
import tempfile


class Table(object):

    '''Represent tabular data for formatting purposes.'''
    
    sep = '  '
    
    def __init__(self):
        self.caption = None
        self.columns = []
        self.rows = []

    def add_column(self, heading1, heading2, format, left=False):
        self.columns.append((heading1, heading2, format, left))

    def add_row(self, data):
        assert len(data) == len(self.columns)
        self.rows.append(data)

    def write_plaintext(self, f):
        if self.caption:
            f.write('%s\n%s\n\n' % (self.caption, '-' * len(self.caption)))
        
        cells = []
        cells.append([h1 for h1, h2, format, left in self.columns])
        cells.append([h2 for h1, h2, format, left in self.columns])
        for row in self.rows:
            cells.append([self.format_cell(row[i], self.columns[i][2])
                          for i in range(len(self.columns))])

        widths = self.compute_column_widths(cells)
        
        f.write('%s\n' % self.format_headings(widths, 0))
        f.write('%s\n' % self.format_headings(widths, 1))
        for row in self.rows:
            f.write('%s\n' % self.format_row(row, widths))

    def format_cell(self, data, format):
        return format % data

    def compute_column_widths(self, cells):
        widths = [0] * len(self.columns)
        for row in cells:
            for i, data in enumerate(row):
                widths[i] = max(widths[i], len(data))
        return widths

    def format_headings(self, widths, which):
        headings = [self.pad(widths[i], 
                             self.columns[i][which], 
                             self.columns[i][3])
                    for i in range(len(widths))]
        return self.sep.join(headings)
        
    def format_row(self, row, widths):
        def cell(i):
            h1, h2, format, left = self.columns[i]
            return self.pad(widths[i], format % row[i], left)
        cells = [cell(i) for i in range(len(widths))]
        return self.sep.join(cells)

    def pad(self, width, text, left):
        if left:
            return '%-*s' % (width, text)
        else:
            return '%*s' % (width, text)


class Benchmarker(cliapp.Application):

    def add_settings(self):
        self.settings.string_list(['command'], 'command to benchmark')
        self.settings.string_list(['verify'], 'command for verifying result')
        self.settings.string_list(['setup'], 'command for setup')
        self.settings.string_list(['cleanup'], 'command for cleanup')
        self.settings.boolean(['setup-once'], 'do setups only once')
        self.settings.boolean(['cleanup-once'], 'do cleanups only once')
        self.settings.boolean(['verbose'], 'show commands and their output')
        
    def process_args(self, args):
        results = Table()
        results.add_column('user', '(s)', '%.1f')
        results.add_column('system', '(s)', '%.1f')
        results.add_column('real', '(s)', '%.1f')
        results.add_column('max RSS', '(KiB)', '%d')
        results.add_column('cmd', '', '%-0s', left=True)

        commands = self.settings['command']
        if commands and self.settings['setup-once']:
            self.setup()
        for cmd in commands:
            if not self.settings['setup-once']:
                self.setup()
            numbers = self.measure_cmd(cmd)
            self.verify()
            results.add_row(numbers + (cmd,))
            if not self.settings['cleanup-once']:
                self.cleanup()
            if self.settings['verbose']:
                self.output.write('\n')
        if commands and self.settings['cleanup-once']:
            self.cleanup()

        results.write_plaintext(self.output)

    def setup(self):
        self.run_shell(self.settings['setup'])

    def cleanup(self):
        self.run_shell(self.settings['cleanup'])

    def verify(self):
        self.run_shell(self.settings['verify'])

    def run_shell(self, cmds):
        for cmd in cmds:
            if self.settings['verbose']:
                self.output.write('COMMAND: %s\n' % cmd)
            out = self.runcmd(['sh', '-c', cmd])
            if self.settings['verbose']:
                self.output.write(out)

    def measure_cmd(self, cmd):
        fd, timings = tempfile.mkstemp()
        time_argv = ['/usr/bin/time', 
                    '-o', timings, 
                     '--format', '%U\n%S\n%e\n%M\n%e',
                     '--',
                     'sh',
                     '-c', cmd]

        if self.settings['verbose']:
            self.output.write('MEASURE: %s\n' % cmd)
        out = self.runcmd(time_argv)
        if self.settings['verbose']:
            self.output.write(out)

        data = []
        while True:
            datum = os.read(fd, 1024**2)
            if not datum:
                break
            data.append(datum)
        os.close(fd)
        data = ''.join(data)

        fields = [float(x) for x in data.splitlines()]
        user = fields[0]
        system = fields[1]
        real = fields[2]
        maxrss = fields[3]
        elapsed = fields[4]
        return user, system, real, maxrss


Benchmarker().run()