#!/usr/bin/python
# Copyright (C) 2012-2015  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, see <http://www.gnu.org/licenses/>.


'''A Morph deployment write extension for raw disk images.'''


import contextlib
import os
import pyfdisk
import re
import subprocess
import sys
import time
import tempfile

import writeexts


class RawDiskWriteExtension(writeexts.WriteExtension):

    '''See rawdisk.write.help for documentation'''

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

        temp_root, location = args
        upgrade = self.get_environment_boolean('UPGRADE')

        if upgrade:
            self.upgrade_local_system(location, temp_root)
        else:
            try:
                if not self.is_device(location):
                    with self.created_disk_image(location):
                        self.create_baserock_system(temp_root, location)
                    self.status(msg='Disk image has been created at %s' %
                                     location)
                else:
                    self.create_baserock_system(temp_root, location)
                    self.status(msg='System deployed to %s' % location)
            except Exception:
                self.status(msg='Failure to deploy system to %s' %
                                 location)
                raise

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

        try:
            with self.mount(location) as mp:
                self.do_upgrade(mp, temp_root)
                return
        except subprocess.CalledProcessError:
            pass

        # At this point, we have failed to mount a raw image, so instead
        # search for a Baserock root filesystem in the device's partitions
        with self.find_and_mount_rootfs(location) as mp:
            self.do_upgrade(mp, temp_root)

    def do_upgrade(self, mp, temp_root):
        version_label = self.get_version_label(mp)
        self.status(msg='Updating image to a new version with label %s' %
                    version_label)

        version_root = os.path.join(mp, 'systems', version_label)
        os.mkdir(version_root)

        old_orig = os.path.join(mp, 'systems', 'default', 'orig')
        new_orig = os.path.join(version_root, 'orig')
        subprocess.check_call(
            ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])

        subprocess.check_call(
            ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
             temp_root + os.path.sep, new_orig])

        self.create_run(version_root)

        default_path = os.path.join(mp, 'systems', 'default')
        if os.path.exists(default_path):
            os.remove(default_path)
        else:
            # we are upgrading and old system that does
            # not have an updated extlinux config file
            if self.bootloader_config_is_wanted():
                self.generate_bootloader_config(mp)
            self.install_bootloader(mp)
        os.symlink(version_label, default_path)

        if self.bootloader_config_is_wanted():
            self.install_kernel(version_root, temp_root)

    def get_version_label(self, mp):
        version_label = os.environ.get('VERSION_LABEL')

        if version_label is None:
            raise writeexts.ExtensionError('VERSION_LABEL was not given')

        if os.path.exists(os.path.join(mp, 'systems', version_label)):
            raise writeexts.ExtensionError('VERSION_LABEL %s already exists'
                                           % version_label)

        return version_label


RawDiskWriteExtension().run()