#!/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):
    ''' Evaluates 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']) > StrictVersion(vers):
        return False
    return True

class VersionError(Exception):
    pass

class Upgrade(object):
    def __init__(self, ur_dir, to_version, mirror=DEFAULT_MIRROR, **kwargs):
        self.mirror = mirror
        self.ureka = Ureka(ur_dir)
        self.ureka_next = None
        self.to_version = to_version
        self.tmp = tempfile.mkdtemp(prefix="upgrade")
        self.tmp_dist = os.path.join(self.tmp, self.to_version)
        self.archive_ext = '.tar.gz'

        if 'archive_ext' in kwargs:
            self.archive_ext = kwargs['archive_ext']

        self.archive = "Ureka_{}_{}_{}{}".format(self.ureka['os'],
                                                 self.ureka['bits'],
                                                 self.to_version,
                                                 self.archive_ext)
        self.archive_url = "{}/{}/{}".format(self.mirror,
                                             self.to_version,
                                             self.archive)
        self.archive_path = os.path.abspath(os.path.join(self.tmp, self.archive))

        if not ur_check_version(self.ureka, self.to_version):
            self._cleanup()
            raise VersionError("Refusing to downgrade from {} to {}".format(self.ureka['version'], self.to_version))
            exit(1)

        signal.signal(signal.SIGINT, self._cleanup_on_signal)
        signal.signal(signal.SIGTERM, self._cleanup_on_signal)

        if not which('rsync'):
            print('++ rsync not found in PATH. Please install it.')
            exit(1)

    def run(self):
        self._get_archive()
        ureka_next_path = self._unpack_archive()
        self._pre()

        # Populate temporary Ureka upgrade object
        self.ureka_next = Ureka(ureka_next_path)

        # Sync data
        if not self._upgrade(self.ureka_next, self.ureka):
            self._cleanup()
            print("++ Upgrade failed!")
            return 1

        # Regenerate original Ureka object
        self.ureka = Ureka(self.ureka.path)

        if not self._post(self.ureka):
            self._cleanup()
            print("++ Post-installation failed!")
            return 1

        self._cleanup()

    def _cleanup(self):
        shutil.rmtree(self.tmp)

    def _cleanup_on_signal(self, sig, stack):
        print('')
        print('++ Received signal {}...'.format(sig))
        self._cleanup()
        exit(sig)

    def _get_archive(self):
        print("+ Downloading archive... {}".format(self.archive_url))
        try:
            req_data = urllib2.urlopen(self.archive_url)
        except Exception as ex:
            print("Error {}, {}: {}".format(ex.code, ex.msg, self.archive_url))
            exit(1)

        remote_size = float(req_data.headers['content-length'])

        total = 0
        with open(self.archive_path, 'wb') as installer:
            for chunk in iter(lambda: req_data.read(16 * 1024), ''):
                total += float(len(chunk))
                installer.write(chunk)
                print("\r{:.2f}% [{:.2f}/{:.2f} MB]".format((total / remote_size * 100),
                                                            (total / 1024 ** 2),
                                                            (remote_size / 1024 ** 2))),
        print("")

    def _unpack_archive(self):
        print('+ Preparing to unpack (please wait)...'),
        installer = tarfile.open(self.archive_path, 'r')
        files_total = len(installer.getmembers())
        print('done')
        print("+ Unpacking...")
        for index, member in enumerate(installer, start=1):
            print("\r{:.2f}% [{}/{}]".format((index / files_total * 100),
                                             index,
                                             files_total)),
            installer.extract(member, path=self.tmp_dist)
        print("")

        return os.path.join(self.tmp_dist, 'Ureka')

    def _pre(self):
        print("+ Executing pre-upgrade tasks...")
        misc = os.path.join(self.tmp_dist, 'Ureka', 'misc', 'name')
        with open(misc, 'w+') as fp:
            fp.write(self.ureka['name'] + os.linesep)

    def _upgrade(self, src, dest):
        print("+ Upgrade in progress ({} to {})...".format(src['version'],
                                                               dest['version']))
        command = 'rsync -a -u {} {}'.format(dest.path,
                                             os.path.dirname(src.path)).split()
        proc = subprocess.Popen(command)
        proc.wait()

        if proc.returncode:
            return False
        return True

    def _post(self, ur):
        print("+ Executing post-upgrade tasks...")
        ur_env = ur_getenv(ur.path)

        command = 'ur_normalize -n -i -x'.split()
        proc = subprocess.Popen(command, env=ur_env)
        proc.wait()

        if proc.returncode:
            return False
        return True


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)

    mirror = DEFAULT_MIRROR
    if args.mirror:
        mirror = args.mirror

    upgrade = Upgrade(args.UR_DIR, args.request, mirror)
    exit_code = upgrade.run()

    exit(exit_code)