#!/usr/bin/python
#
# Copyright 2010, 2011  Lars Wirzenius
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import cliapp
import ConfigParser
import glob
import logging
import os
import shutil
import socket
import subprocess
import tempfile


class ObnamBenchmark(cliapp.Application):

    default_sizes = ['1g/100m']
    keyid = '3B1802F81B321347'
    opers = ('backup', 'restore', 'list_files', 'forget')

    def add_settings(self):
        self.settings.string(['results'], 'put results under DIR (%default)',
                            metavar='DIR', default='../benchmarks')
        self.settings.string(['obnam-branch'],
                             'use DIR as the obnam branch to benchmark '
                                '(default: %default)',
                              metavar='DIR',
                              default='.')
        self.settings.string(['larch-branch'],
                             'use DIR as the larch branch (default: %default)',
                             metavar='DIR',
                            )
        self.settings.string(['seivot-branch'],
                             'use DIR as the seivot branch '
                                '(default: installed seivot)',
                             metavar='DIR')
        self.settings.boolean(['with-encryption'],
                              'run benchmark using encryption')

        self.settings.string(['profile-name'],
                             'short name for benchmark scenario',
                             default='unknown')
        self.settings.string_list(['size'],
                                  'add PAIR to list of sizes to '
                                    'benchmark (e.g., 10g/1m)',
                                  metavar='PAIR')
        self.settings.bytesize(['file-size'], 'how big should files be?',
                               default=4096)
        self.settings.integer(['generations'],
                              'benchmark N generations (default: %default)',
                              metavar='N',
                              default=5)
        self.settings.boolean(['use-sftp-repository'],
                              'access the repository over SFTP '
                                '(requires ssh to localhost to work)')
        self.settings.boolean(['use-sftp-root'],
                              'access the live data over SFTP '
                                '(requires ssh to localhost to work)')
        self.settings.integer(['sftp-delay'],
                              'add artifical delay to sftp transfers '
                                '(in milliseconds)')
        self.settings.string(['description'], 'describe benchmark')
        self.settings.boolean(['drop-caches'], 'drop kernel buffer caches')
        self.settings.string(['seivot-log'], 'seivot log setting')

        self.settings.boolean(['verify'], 'verify restores')

    def process_args(self, args):
        self.require_tmpdir()

        obnam_revno = self.bzr_revno(self.settings['obnam-branch'])
        if self.settings['larch-branch']:
            larch_revno = self.bzr_revno(self.settings['larch-branch'])
        else:
            larch_revno = None

        results = self.results_dir(obnam_revno, larch_revno)

        obnam_branch = self.settings['obnam-branch']
        if self.settings['seivot-branch']:
            seivot = os.path.join(self.settings['seivot-branch'], 'seivot')
        else:
            seivot = 'seivot'

        generations = self.settings['generations']

        tempdir = tempfile.mkdtemp()
        env = self.setup_gnupghome(tempdir)

        sizes = self.settings['size'] or self.default_sizes
        logging.debug('sizes: %s' % repr(sizes))

        file_size = self.settings['file-size']
        profile_name = self.settings['profile-name']

        for pair in sizes:
            initial, inc = self.parse_size_pair(pair)

            msg = 'Profile %s, size %s inc %s' % (profile_name, initial, inc)
            print
            print msg
            print '-' * len(msg)
            print

            obnam_profile = os.path.join(results,
                                         'obnam--%(op)s-%(gen)s.prof')
            output = os.path.join(results, 'obnam.seivot')
            if os.path.exists(output):
                print ('%s already exists, not re-running benchmark' %
                        output)
            else:
                argv = [seivot,
                        '--obnam-branch', obnam_branch,
                        '--incremental-data', inc,
                        '--file-size', str(file_size),
                        '--obnam-profile', obnam_profile,
                        '--generations', str(generations),
                        '--profile-name', profile_name,
                        '--sftp-delay', str(self.settings['sftp-delay']),
                        '--initial-data', initial,
                        '--output', output]
                if self.settings['larch-branch']:
                    argv.extend(['--larch-branch', self.settings['larch-branch']])
                if self.settings['seivot-log']:
                    argv.extend(['--log', self.settings['seivot-log']])
                if self.settings['drop-caches']:
                    argv.append('--drop-caches')
                if self.settings['use-sftp-repository']:
                    argv.append('--use-sftp-repository')
                if self.settings['use-sftp-root']:
                    argv.append('--use-sftp-root')
                if self.settings['with-encryption']:
                    argv.extend(['--encrypt-with', self.keyid])
                if self.settings['description']:
                    argv.extend(['--description',
                                 self.settings['description']])
                if self.settings['verify']:
                    argv.append('--verify')
                self.runcmd(argv, env=env)

        shutil.rmtree(tempdir)

    def require_tmpdir(self):
        if 'TMPDIR' not in os.environ:
            raise cliapp.AppException('TMPDIR is not set. '
                                       'You would probably run out of space '
                                       'on /tmp.')
        if not os.path.exists(os.environ['TMPDIR']):
            raise cliapp.AppException('TMPDIR points at a non-existent '
                                        'directory %s' % os.environ['TMPDIR'])
        logging.debug('TMPDIR=%s' % repr(os.environ['TMPDIR']))

    @property
    def hostname(self):
        return socket.gethostname()

    @property
    def obnam_branch_name(self):
        obnam_branch = os.path.abspath(self.settings['obnam-branch'])
        return os.path.basename(obnam_branch)

    def results_dir(self, obnam_revno, larch_revno):
        parent = self.settings['results']
        parts = [self.hostname, self.obnam_branch_name, str(obnam_revno)]
        if larch_revno:
            parts.append(str(larch_revno))
        prefix = os.path.join(parent, "-".join(parts))

        get_path = lambda counter: "%s-%d" % (prefix, counter)

        counter = 0
        dirname = get_path(counter)
        while os.path.exists(dirname):
            counter += 1
            dirname = get_path(counter)
        os.makedirs(dirname)
        return dirname

    def setup_gnupghome(self, tempdir):
        gnupghome = os.path.join(tempdir, 'gnupghome')
        shutil.copytree('test-gpghome', gnupghome)
        env = dict(os.environ)
        env['GNUPGHOME'] = gnupghome
        return env

    def bzr_revno(self, branch):
        p = subprocess.Popen(['bzr', 'revno'], cwd=branch,
                             stdout=subprocess.PIPE)
        out, err = p.communicate()
        if p.returncode != 0:
            raise cliapp.AppException('bzr failed')

        revno = out.strip()
        logging.debug('bzr branch %s has revno %s' % (branch, revno))
        return revno

    def parse_size_pair(self, pair):
        return pair.split('/', 1)


if __name__ == '__main__':
    ObnamBenchmark().run()

