#!/usr/bin/env python
#Copyright (c) 2014, Joseph Hunkeler <jhunkeler at gmail.com>
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer. 
#2. Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from __future__ import division

import argparse
import os
import shutil
import subprocess
import tempfile
import urllib2
import signal
import tarfile
from distutils.version import StrictVersion
from distutils.spawn import find_executable as which
from collections import namedtuple
from string import Template
from itertools import chain

DEFAULT_MIRROR = "http://ssb.stsci.edu/ureka"

class Ureka(object):
    ''' 
    '''
    def __init__(self, basepath):
        self.path = basepath
        self.path_data = os.path.abspath(os.path.join(self.path, 'misc'))
        self._files = ['os', 'bits', 'version', 'name']
        if os.path.exists(self.path_data):
            self.info = self._info()
        else:
            self.info = {}
        

    def _info(self):
        data = []
        for key in self._files:
            path = os.path.join(self.path_data, key)
            item = None
            try:
                item = open(path, 'r').readline().strip()
            except:
                item = ''
            data.append(item)
            
        ureka_info = namedtuple('ureka_info', self._files)
        return ureka_info._make(data)
    
    def __getitem__(self, key):
        info = self.info._asdict()
        if key not in info:
            return None
        return info[key]
    
    def __iter__(self):
        for item in self.info._asdict().iteritems():
            yield item

def ur_getenv(ur_dir):
    ''' Evaluate environment variables produced by ur-setup-real
    '''
    path = os.path.join(ur_dir, 'bin')
    command = '{}/ur-setup-real -sh'.format(path).split()
    output = subprocess.check_output(command, shell=True, stderr=open(os.path.devnull))
    output = output.split(os.linesep)
    output_env = []
    
    # Generate environment keypairs
    for line in output:
        if not line:
            continue
        if line.startswith('export'):
            continue
        line = line.strip()
        line = line.replace(' ', '')
        line = line.replace(';', '')
        line = line.replace('"', '')
        output_env.append(line.partition("=")[::2])
    
    output_env_temp = dict(output_env)
    
    # Perform shell expansion of Ureka's environment variables
    for k, v in output_env_temp.items():
        template = Template(v)
        v = template.safe_substitute(output_env_temp)
        output_env_temp[k] = v
    
    # Merge Ureka's environment with existing system environment
    output_env_temp = dict(chain(os.environ.items(), output_env_temp.items()))
    
    # Assign expanded variables
    output_env = output_env_temp
    
    return output_env

def ur_check_version(urobj, vers):
    if StrictVersion(urobj['version']) > vers:
        return False
    return True

def get_archive(url):
    print("Downloading archive... {}".format(url))
    try:
        req_data = urllib2.urlopen(url)
    except Exception as ex:
        print("Error {}, {}: {}".format(ex.code, ex.msg, url))
        exit(1)
        
    remote_size = float(req_data.headers['content-length'])
    total = 0
    with open(os.path.join(UPGRADE_TEMP, UREKA_TARBALL), 'wb') as installer:
        for chunk in iter(lambda: req_data.read(16*1024), ''):
            print("\r{:.2f}% [{:.2f}/{:.2f} MB]".format((total / remote_size * 100), (total / 1024 ** 2), (remote_size / 1024 ** 2))),
            total += float(len(chunk))
            installer.write(chunk)
    print("")

def upgrade(srcobj, destobj):
    print("Upgrading distribution from {} to {}...".format(srcobj['version'], destobj['version']))
    proc = subprocess.Popen('rsync -a -ureka_dist_original {} {}'.format(destobj.path, os.path.dirname(srcobj.path)).split())
    proc.wait()

def unpack_archive():
    installer = tarfile.open(os.path.join(UPGRADE_TEMP, UREKA_TARBALL), 'r')
    print("Preparing to unpack upgrade... (please wait)")
    files_total = len(installer.getmembers()) 
    
    print("Unpacking upgrade...")
    for index, member in enumerate(installer, start=1):
        print("\r{:.2f}% [{}/{}]".format((index / files_total * 100), index, files_total)),
        installer.extract(member, path=TEMP_INSTALL_PATH)
    print("")

def pre_upgrade():    
    with open(os.path.join(TEMP_INSTALL_PATH, 'Ureka', 'misc', 'name'), 'w+') as fp:
        fp.write(ureka_dist_original['name'] + os.linesep)

def post_upgrade(urobj):
    print("Performing upgrade maintenance tasks...")
    ur_env = ur_getenv(urobj.path)
    proc = subprocess.Popen('ur_normalize -n -i -x'.split(), env=ur_env)
    proc.wait()
    
    print("Purging temporary data...")
    shutil.rmtree(UPGRADE_TEMP)


def cleanup(sig, stack):
    print('')
    print('Received signal {}!'.format(sig))
    print('Cleaning up...')
    shutil.rmtree(UPGRADE_TEMP)
    exit(sig)
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Ureka Upgrade Utility')
    parser.add_argument('UR_DIR', action='store', type=str, help='Absolute path to Ureka installation')
    parser.add_argument('--latest', action='store_true', default=False, help='Upgrade to the latest (stable) version')
    parser.add_argument('--request', type=str, help='Upgrade to a specific version')
    parser.add_argument('--mirror', type=str, help='Use a Ureka download mirror')
    args = parser.parse_args()
    
    if args.latest and args.request:
        print("--latest and --request are mutually exclusive options.")
        exit(1)
    
    if args.latest:
        # Need to work out a deal with SSB.
        # "LATEST-STABLE" and "LATEST-DEVEL" links could be useful 
        print("--latest is not implemented yet.")
        exit(1)
        
    if not args.request:
        print("No version requested. Use --request (e.g. 1.4.1)")
        exit(1)
    
    UPGRADE_TEMP = tempfile.mkdtemp(prefix="upgrade")
    CURRENT = args.UR_DIR
    
    # Register signal handlers AFTER temporary directory is created
    signal.signal(signal.SIGINT, cleanup)
    signal.signal(signal.SIGTERM, cleanup)
             
    if not which('rsync'):
        print('rsync not found in PATH. Please install it.')
        exit(1)
    
    ureka_dist_original = Ureka(CURRENT)
    
    if not ur_check_version(ureka_dist_original, args.request):
        print("Refusing to downgrade from {} to {}".format(ureka_dist_original['version'], args.request))
        exit(1)
    
    TEMP_INSTALL_PATH = os.path.join(UPGRADE_TEMP, args.request)
    UREKA_TARBALL = "Ureka_{}_{}_{}.tar.gz".format(ureka_dist_original['os'], ureka_dist_original['bits'], args.request)

    url = "{}/{}/{}".format(DEFAULT_MIRROR, args.request, UREKA_TARBALL)
    if args.mirror:
        url = "{}/{}/{}".format(args.mirror, args.request, UREKA_TARBALL)

    # Download requested archive
    get_archive(url)
    
    # Unpack archive to temporary storage
    unpack_archive()
    
    # Generate upgrade Ureka object
    ureka_dist_new = Ureka(os.path.join(TEMP_INSTALL_PATH, 'Ureka'))
    
    # Perform tasks required before upgrading
    pre_upgrade()
    
    # Perform upgrade
    upgrade(ureka_dist_original, ureka_dist_new)
    
    # Regenerate upgraded ureka installation target
    ureka_dist_original = Ureka(CURRENT)

    # Perform tasks required after upgrade is completed
    post_upgrade(ureka_dist_original)