#!/usr/bin/python

import argparse
import sys
import os
import subprocess
import zipfile
import guestfs
import tempfile
import tarfile

parser = argparse.ArgumentParser()
parser.add_argument('--image', dest='image', default=None)
parser.add_argument('--root', dest='root', default=None)
parser.add_argument('--partition', dest='partition', default=None)
parser.add_argument('--system32', dest='system32', choices=['Windows/System32', 'WINDOWS/system32'])
parser.add_argument('--erase_old_driver_files', dest='erase', default=False, action='store_true')
parser.add_argument('--allow-old', action='store_true', default=False)
parser.add_argument('--brew')
parser.add_argument('--arch', choices=['x86', 'x64'])
parser.add_argument('--winver', choices=['w7', 'xp'])
args, rest = parser.parse_known_args(sys.argv[1:])

if os.system('which guestfish'):
    print "missing guestfish - yum install guestfish (or something)"
    sys.exit(1)

def no_flags_usage(rest):
    if len(rest) != 4:
        parser.print_usage()
        sys.exit(1)

# do this via argparse??
if not args.image:
    no_flags_usage(rest)
    args.image = rest[0]
if not args.root and not args.brew:
    no_flags_usage(rest)
    args.root = rest[1]
if not args.partition and not args.winver:
    no_flags_usage(rest)
    args.partition = rest[2]
if not args.system32:
    args.system32 = 'WINDOWS/system32' if (len(rest) != 4) else rest[3]
if args.winver:
    # override system32 logic before
    args.system32 = 'Windows/System32' if args.winver == 'w7' else 'WINDOWS/system32'
    if not args.partition:
        args.partition = 2 if args.winver == 'w7' else 1

HOME = os.environ['HOME']
BREW_CACHE='%s/.brew_qxl_win_cache' % HOME

class chdir:
    """ usage:
    with chdir('/tmp'):
        print os.getcwd() # prints '/tmp'
    """
    def __init__(self, d):
        self.d = d
        if not os.path.exists(d):
            raise Exception('chdir: no directory %r' % d)

    def __enter__(self, *args, **kw):
        self.org = os.getcwd()
        os.chdir(self.d)

    def __exit__(self, *args, **kw):
        os.chdir(self.org)

# curl invocation notes:
#  -R/--remote-time
#      When used, this will make libcurl attempt to figure out the timestamp of the remote file, and if that is available make the local file get that same timestamp.

def curl(url):
    print "downloading %s" % url
    os.system('curl -R -O %s' % url)

def fill_brew_driver_cache(ver, winver, arch, debug=False):
    # instead of using the brew utility. can't grok download-latest
    assert(winver in ['w7', 'xp'])
    assert((winver == 'xp' and arch == 'x86') or arch in ['x86', 'x64'])
    driver_path = os.path.join(BREW_CACHE, os.path.join(*ver.split('-')))
    if not os.path.exists(driver_path):
        os.makedirs(driver_path)
    print "ver = %r" % ver
    base_url = 'http://download.devel.redhat.com/brewroot/packages/qxl-win/%s/%s/win/' % tuple(ver.split('-'))
    saved_wd = os.getcwd()
    with chdir(driver_path):
        for mid in ['w7_x86', 'w7_x64', 'xp_x86']:
            for post in ['', '_debug']:
                filename = 'qxl_%s%s.zip' % (mid, post)
                target = os.path.join(driver_path, filename)
                if not os.path.exists(target):
                    curl(base_url + filename)
                if not os.path.exists(target):
                    print "curl failed?"
                    sys.exit(1)
    return driver_path + '/qxl_%s_%s%s.zip' % (winver, arch, '' if not debug else '_debug')

def install_brew_driver(args):
    """
    get a driver from brew and install in an image
    args.brew - version (0.1-8)
    args.winver - windows version. TODO - extract from image.
    args.arch - architecture. TODO - extract from image

    Unpacks the driver files under the brew cache directory.
    TODO - also grab the pdbs?
    """
    if not args.winver or not args.arch:
        print "missing --arch and --winver"
        sys.exit(1)
    winver = args.winver
    arch = args.arch
    zip_filename = fill_brew_driver_cache(ver=args.brew, winver=winver, arch=arch)
    zip_dirname = os.path.dirname(zip_filename)
    args.root = os.path.join(zip_dirname, args.winver, args.arch)
    if not os.path.exists(args.root):
        print "unpacking %s" % zip_filename
        with chdir(zip_dirname):
            zipfile.ZipFile(zip_filename).extractall()
    install_driver(args)

def install_driver(args):
    base=os.path.dirname(args.image)
    qxldd_path = os.path.realpath(os.path.join(args.root, 'qxldd.dll'))
    if not os.path.exists(qxldd_path):
        print "missing qxldd.dll"
        sys.exit(1)
    qxlsys_path = os.path.realpath(os.path.join(args.root, 'qxl.sys'))
    if not os.path.exists(qxlsys_path):
        print "missing qxl.sys"
        sys.exit(1)
    if not os.path.exists(args.image):
        print "provided image does not exist"
        sys.exit(1)
    image_user = os.popen("lsof -t %s" % args.image).read().strip()
    if len(image_user) != 0:
        print "image is being used by %s" % image_user
        print "will not update used image"
        sys.exit(1)
    not_too_old(qxldd_path, not args.allow_old)
    not_too_old(qxlsys_path, not args.allow_old)

    print "copying qxl driver %s to %s" % (args.root, args.image)
    g = guestfs.GuestFS()
    g.set_trace(1)
    g.add_drive(args.image)
    g.launch()
    g.mount("/dev/vda%s" % args.partition, "/")
    for local_fname, target_fname in [
        (qxldd_path, '/' + args.system32 + '/qxldd.dll'),
        (qxlsys_path, '/' + args.system32 + '/qxlsys.dll')]:
        g.upload(local_fname, target_fname)
        g.checksum('md5', target_fname)
    temp_dir = tempfile.mkdtemp()
    tar_name = os.path.join(temp_dir, 'qxl.tar')
    with chdir(args.root):
        tar = tarfile.open(tar_name, 'w')
        tar.add('.')
        tar.close()
    if not g.exists('/qxl'):
        g.mkdir('/qxl')
    g.tar_in(tar_name, '/qxl/')
    os.system("md5sum %s" % qxldd_path)
    os.system("md5sum %s" % qxlsys_path)

def not_too_old(p, do_exit):
    if os.system("peversion %s" % p):
        print "%s is too old" % p
        if do_exit:
            sys.exit(1)

def guestfish(commands):
    gf = subprocess.Popen('guestfish', stdin=subprocess.PIPE)
    gf.stdin.write(commands)
    gf.stdin.close()
    return gf.wait()

if __name__ == '__main__':
    # get driver from brew
    if args.brew:
        args.allow_old = True
        install_brew_driver(args)
    else:
        install_driver(args)