#!/usr/bin/python3
#
# This file is part of Plinth.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Configuration helper for disks manager.
"""

import argparse
import re
import subprocess
import sys


def parse_arguments():
    """Return parsed command line arguments as dictionary."""
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='subcommand', help='Sub command')

    subparser = subparsers.add_parser(
        'is-partition-expandable',
        help='Return whether a given partition can be expanded')
    subparser.add_argument(
        'device', help='Partition for which check needs to be performed')

    subparser = subparsers.add_parser(
        'expand-partition',
        help='Expand a partition to take adjacent free space')
    subparser.add_argument(
        'device', help='Partition which needs to be resized')

    return parser.parse_args()


def subcommand_is_partition_expandable(arguments):
    """Return a list of partitions that can be expanded."""
    _, _, free_space = _get_free_space(arguments.device)
    print(free_space['size'])


def subcommand_expand_partition(arguments):
    """Expand a partition to take adjacent free space."""
    device = arguments.device
    device, requested_partition, free_space = _get_free_space(device)

    if requested_partition['table_type'] == 'msdos' and \
       int(requested_partition['number']) >= 5:
        print('Expanding logical partitions currently unsupported',
              file=sys.stderr)
        sys.exit(4)

    _resize_partition(device, requested_partition, free_space)
    _resize_file_system(device, requested_partition, free_space)


def _resize_partition(device, requested_partition, free_space):
    """Resize the partition table entry."""
    command = ['parted', '--align=optimal', '--script', device, 'unit', 'B',
               'resizepart', requested_partition['number'],
               str(free_space['end'])]
    try:
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as exception:
        print('Error expanding partition:', exception, file=sys.stderr)
        sys.exit(5)


def _resize_file_system(device, requested_partition, free_space):
    """Resize a file system inside a partition."""
    if requested_partition['type'] == 'btrfs':
        _resize_btrfs(device, requested_partition, free_space)
    elif requested_partition['type'] == 'ext4':
        _resize_ext4(device, requested_partition, free_space)


def _resize_ext4(device, requested_partition, free_space):
    """Resize an ext4 file system inside a partition."""
    partition_device = _get_partition_device(
        device, requested_partition['number'])
    try:
        command = ['resize2fs', partition_device]
        subprocess.run(command, stdout=subprocess.DEVNULL,
                       stderr=subprocess.DEVNULL, check=True)
    except subprocess.CalledProcessError as exception:
        print('Error expanding filesystem:', exception, file=sys.stderr)
        sys.exit(6)


def _resize_btrfs(device, requested_partition, free_space):
    """Resize a btrfs file system inside a partition."""
    try:
        command = ['btrfs', 'filesystem', 'resize', 'max', '/']
        subprocess.run(command, stdout=subprocess.DEVNULL, check=True)
    except subprocess.CalledProcessError as exception:
        print('Error expanding filesystem:', exception, file=sys.stderr)
        sys.exit(6)


def _get_free_space(device):
    """Return the amount of free space after a partition."""
    device, partition_number = \
        _get_root_device_and_partition_number(device)

    try:
        requested_partition, free_spaces = \
            _get_partitions_and_free_spaces(device, partition_number)
    except Exception as exception:
        print('Error getting partition details: ', exception, file=sys.stderr)
        sys.exit(2)

    # Don't accept extended partitions for now
    if requested_partition['table_type'] == 'msdos' and \
       int(requested_partition['number']) >= 5:
        print('Expanding logical partitions currently unsupported',
              file=sys.stderr)
        sys.exit(3)

    # Don't accept anything but btrfs and ext4 filesystems
    if requested_partition['type'] not in ('btrfs', 'ext4'):
        print('Unsupported file system type: ', requested_partition['type'],
              file=sys.stderr)
        sys.exit(4)

    found_free_space = None
    for free_space in free_spaces:
        if free_space['start'] != requested_partition['end'] + 1:
            continue

        if free_space['size'] < 10 * 1024 * 1024:  # Minimum 10MiB
            continue

        found_free_space = free_space

    if not found_free_space:
        sys.exit(5)

    return device, requested_partition, found_free_space


def _get_partition_device(device, partition_number):
    """Return the device corresponding to a parition in a given device."""
    if re.match('[0-9]', device[-1]):
        return device + 'p' + str(partition_number)

    return device + str(partition_number)


def _get_root_device_and_partition_number(device):
    """Return the parent device and number of partition separately."""
    match = re.match('(.+[a-zA-Z]\d+)p(\d+)$', device)
    if not match:
        match = re.match('(.+[a-zA-Z])(\d+)$', device)
        if not match:
            print('Invalid device, must be a partition', file=sys.stderr)
            sys.exit(1)

    return match.group(1), match.group(2)


def _get_partitions_and_free_spaces(device, partition_number):
    """Run parted and return list of partitions and free spaces."""
    command = ['parted', '--machine', '--script', device, 'unit', 'B',
               'print', 'free']
    process = subprocess.run(command, stdout=subprocess.PIPE, check=True)

    requested_partition = None
    free_spaces = []

    lines = process.stdout.decode().splitlines()
    partition_table_type = lines[1].split(':')[5]
    for line in lines[2:]:
        line = line.rstrip(';')
        keys = ('number', 'start', 'end', 'size', 'type')
        parts = line.split(':')
        segment = dict(zip(keys, parts[:5]))

        segment['table_type'] = partition_table_type
        segment['start'] = _interpret_unit(segment['start'])
        segment['end'] = _interpret_unit(segment['end'])
        segment['size'] = _interpret_unit(segment['size'])

        if segment['type'] == 'free':
            segment['number'] = None
            free_spaces.append(segment)
        else:
            if segment['number'] == partition_number:
                requested_partition = segment

    return requested_partition, free_spaces


def _interpret_unit(value):
    """Return value in bytes after understanding parted unit."""
    value = value.rstrip('B')  # For now, we only need to understand bytes
    return int(value)


def main():
    """Parse arguments and perform all duties."""
    arguments = parse_arguments()

    subcommand = arguments.subcommand.replace('-', '_')
    subcommand_method = globals()['subcommand_' + subcommand]
    subcommand_method(arguments)


if __name__ == '__main__':
    main()
