#!/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 module providing Baserock-specific partitioning functions"""

import pyfdisk
import writeexts

def do_partitioning(location, disk_size, temp_root, part_spec):
    '''Perform partitioning

    Perform partitioning using the pyfdisk.py module. Documentation
    for this, and guidance on how to create a partition specification can
    be found in extensions/pyfdisk.README

    This function also validates essential parts of the partition layout

    Args:
        location: Path to the target device or image
        temp_root: Location of the unpacked Baserock rootfs
        part_spec: Path to a YAML formatted partition specification
    Returns:
        A pyfdisk.py Device object
    Raises:
        writeexts.ExtensionError
    '''
    # Create partition table and filesystems
    try:
        dev = pyfdisk.load_yaml(location, disk_size, part_spec)
        writeexts.Extension.status(msg='Loaded partition specification: %s' %
                                        part_spec)

        # FIXME: GPT currently not supported due to missing tools
        if dev.partition_table_format.lower() == 'gpt':
            raise writeexts.ExtensionError('GPT partition tables are not '
                                           'currently supported')

        writeexts.Extension.status(msg=str(dev.partitionlist))
        writeexts.Extension.status(msg='Writing partition table')
        dev.commit()
        dev.create_filesystems(skip=['/'])
    except (pyfdisk.PartitioningError, pyfdisk.FdiskError) as e:
        raise writeexts.ExtensionError(e.msg)

    mountpoints = set(part.mountpoint for part in dev.partitionlist
                                      if hasattr(part, 'mountpoint'))
    if '/' not in mountpoints:
        raise writeexts.ExtensionError('No partition with root '
                                       'mountpoint, please specify a '
                                       'partition with \'mountpoint: /\' '
                                       'in the partition specification')

    mounted_partitions = set(part for part in dev.partitionlist
                             if hasattr(part, 'mountpoint'))

    # Create root filesystem, and copy files to partitions
    for part in mounted_partitions:
        if not hasattr(part, 'filesystem'):
            raise writeexts.ExtensionError('Cannot mount a partition '
                                 'without filesystem, please specify one '
                                 'for \'%s\' partition in the partition '
                                 'specification' % part.mountpoint)
        if part.mountpoint == '/':
            # Check that bootable flag is set for MBR devices
            if (hasattr(part, 'boot')
                  and str(part.boot).lower() not in ('yes', 'true')
                  and dev.partition_table_format.lower() == 'mbr'):
                writeexts.Extension.status(msg='WARNING: Boot partition '
                                               'needs bootable flag set to '
                                               'boot with extlinux/syslinux')

    return dev

def write_raw_files(location, temp_root, dev_or_part, start_offset=0):
    '''Write files with `dd`'''
    offset = start_offset
    for raw_args in dev_or_part.raw_files:
        r = RawFile(temp_root, offset, **raw_args)
        offset = r.next_offset
        r.dd(location)


class RawFile(object):
    '''A class to hold information about a raw file to write to a device'''

    def __init__(self, source_root, wr_offset=0, sector_size=512, **kwargs):
        '''Initialisation function

        Args:
            source_root: Base path for filenames
            wr_offset: Offset to write to (and offset per-file offsets by)
            sector_size: Device sector size (default: 512)
            **kwargs:
                file: A path to the file to write (combined with source_root)
                offset_sectors: An offset to write to in sectors (optional)
                offset_bytes: An offset to write to in bytes (optional)
        '''
        if 'file' not in kwargs:
            raise writeexts.ExtensionError('Missing file name or path')
        self.path = os.path.join(source_root,
                                 re.sub('^/', '', kwargs['file']))

        if not os.path.exists(self.path):
            raise writeexts.ExtensionError('File not found: %s' % self.path)
        elif os.path.isdir(self.path):
            raise writeexts.ExtensionError('Can only dd regular files')
        else:
            self.size = os.stat(self.path).st_size

        self.offset = wr_offset
        if 'offset_bytes' in kwargs:
            self.offset += kwargs['offset_bytes']
        elif 'offset_sectors' in kwargs:
            self.offset += kwargs['offset_sectors'] * sector_size

        # Offset of the first free byte after this file
        self.next_offset = self.size + self.offset

    def dd(self, location):
        writeexts.Extension.status(msg='Writing %s at %d bytes' %
                                       (self.path, self.offset))
        subprocess.check_call(['dd', 'if=%s' % self.path,
                                     'of=%s' % location, 'bs=1',
                                     'seek=%s' % self.offset, 'conv=notrunc'])
        subprocess.check_call('sync')