#!/usr/bin/python
# Copyright (C) 2013-2014  Codethink Limited
#
# 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; version 2 of the License.
#
# 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, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.


'''A Morph deployment write extension for upgrading systems over ssh.'''


import cliapp
import os
import sys
import time
import tempfile

import morphlib.writeexts


def ssh_runcmd_ignore_failure(location, command, **kwargs):
    try:
        return cliapp.ssh_runcmd(location, command, **kwargs)
    except cliapp.AppException:
        pass


class SshRsyncWriteExtension(morphlib.writeexts.WriteExtension):

    '''Upgrade a running baserock system with ssh and rsync.

    It assumes the system is baserock-based and has a btrfs partition.

    The location command line argument is the 'user@hostname' string
    that will be passed to ssh and rsync

    '''

    def process_args(self, args):
        if len(args) != 2:
            raise cliapp.AppException('Wrong number of command line args')

        temp_root, location = args

        self.upgrade_remote_system(location, temp_root)

    def upgrade_remote_system(self, location, temp_root):
        self.complete_fstab_for_btrfs_layout(temp_root)

        root_disk = self.find_root_disk(location)
        version_label = os.environ.get('VERSION_LABEL')
        autostart = self.get_environment_boolean('AUTOSTART')

        self.status(msg='Creating remote mount point')
        remote_mnt = cliapp.ssh_runcmd(location, ['mktemp', '-d']).strip()
        try:
            self.status(msg='Mounting root disk')
            cliapp.ssh_runcmd(location, ['mount', root_disk, remote_mnt])
        except Exception as e:
            ssh_runcmd_ignore_failure(location, ['rmdir', remote_mnt])
            raise e

        try:
            version_root = os.path.join(remote_mnt, 'systems', version_label)
            orig_dir = os.path.join(version_root, 'orig')

            self.status(msg='Creating %s' % version_root)
            cliapp.ssh_runcmd(location, ['mkdir', version_root])

            self.create_remote_orig(location, version_root, remote_mnt,
                                    temp_root)

            # Use the system-version-manager from the new system we just
            # installed, so that we can upgrade from systems that don't have
            # it installed.
            self.status(msg='Calling system-version-manager to deploy upgrade')
            deployment = os.path.join('/systems', version_label, 'orig')
            system_config_sync = os.path.join(
                    remote_mnt, 'systems', version_label, 'orig', 'usr', 'bin',
                    'baserock-system-config-sync')
            system_version_manager = os.path.join(
                    remote_mnt, 'systems', version_label, 'orig', 'usr', 'bin',
                    'system-version-manager')
            cliapp.ssh_runcmd(location,
                    ['env', 'BASEROCK_SYSTEM_CONFIG_SYNC='+system_config_sync,
                     system_version_manager, 'deploy', deployment])

            self.status(msg='Setting %s as the new default system' %
                    version_label)
            cliapp.ssh_runcmd(location,
                    [system_version_manager, 'set-default', version_label])
        except Exception as e:
            self.status(msg='Deployment failed')
            ssh_runcmd_ignore_failure(
                    location, ['btrfs', 'subvolume', 'delete', orig_dir])
            ssh_runcmd_ignore_failure(
                    location, ['rm', '-rf', version_root])
            raise e
        finally:
            self.status(msg='Removing temporary mounts')
            cliapp.ssh_runcmd(location, ['umount', remote_mnt])
            cliapp.ssh_runcmd(location, ['rmdir', remote_mnt])

        if autostart:
            self.status(msg="Rebooting into new system ...")
            ssh_runcmd_ignore_failure(location, ['reboot'])

    def create_remote_orig(self, location, version_root, remote_mnt,
                           temp_root):
        '''Create the subvolume version_root/orig on location'''

        self.status(msg='Creating "orig" subvolume')
        old_orig = self.get_old_orig(location, remote_mnt)
        new_orig = os.path.join(version_root, 'orig')
        cliapp.ssh_runcmd(location, ['btrfs', 'subvolume', 'snapshot',
                                     old_orig, new_orig])

        cliapp.runcmd(['rsync', '-as', '--checksum', '--numeric-ids',
                       '--delete', temp_root + os.path.sep,
                       '%s:%s' % (location, new_orig)])

    def get_old_orig(self, location, remote_mnt):
        '''Identify which subvolume to snapshot from'''

        # rawdisk upgrades use 'factory'
        return os.path.join(remote_mnt, 'systems', 'factory', 'orig')

    def find_root_disk(self, location):
        '''Read /proc/mounts on location to find which device contains "/"'''

        self.status(msg='Finding device that contains "/"')
        contents = cliapp.ssh_runcmd(location, ['cat', '/proc/mounts'])
        for line in contents.splitlines():
            line_words = line.split()
            if (line_words[1] == '/' and line_words[0] != 'rootfs'):
                return line_words[0]


SshRsyncWriteExtension().run()