#
# Copyright (C) 2022 Soleta Networks <info@soleta.eu>
#
# 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.

import logging
import os
import subprocess
import shlex

from subprocess import DEVNULL, PIPE, STDOUT

import psutil

from src.utils.disk import get_partition_device


def find_mountpoint(path):
    """
    Returns mountpoint of a given path
    """
    path = os.path.abspath(path)
    while not os.path.ismount(path):
        path = os.path.dirname(path)
    return path


def mount_mkdir(source, target):
    """
    Mounts and creates the mountpoint directory if it's not present.

    Return True if mount is sucessful or if target is already a mountpoint.
    """
    if not os.path.exists(target):
        os.mkdir(target)

    if not os.path.ismount(target):
        return mount(source, target)
    else:
        return True

    return False


def mount(source, target):
    """
    Mounts source into target directoru using mount(8).

    Return true if exit code is 0. False otherwise.
    """
    cmd = f'mount {source} {target}'
    proc = subprocess.run(cmd.split(), stderr=DEVNULL)

    return not proc.returncode


def umount(target):
    """
    Umounts target using umount(8).

    Return true if exit code is 0. False otherwise.
    """
    cmd = f'umount {target}'
    proc = subprocess.run(cmd.split(), stderr=DEVNULL)

    return not proc.returncode


def umount_all():
    """
    Umounts all mountpoints present in the /mnt folder.
    """
    for path in ['/mnt/'+child for child in os.listdir('/mnt/')]:
        if os.path.ismount(path):
            umount(path)


def get_usedperc(mountpoint):
    """
    Returns percetage of used filesystem as decimal number.
    """
    try:
        total, used, free, perc = psutil.disk_usage(mountpoint)
    except FileNotFoundError:
        return '0'
    return str(perc)


def ogReduceFs(disk, part):
    """
    Shrink filesystem of a partition. Supports ext4 and ntfs partitions.
    Unsupported filesystem or invalid paths don't raise an exception,
    instead this method logs a warning message and does nothing.
    """
    partdev = get_partition_device(disk, part)
    fstype = get_filesystem_type(partdev)

    umount(partdev)
    if fstype == 'ext4':
        _reduce_resize2fs(partdev)
    elif fstype == 'ntfs':
        _reduce_ntfsresize(partdev)
    else:
        logging.warn(f'Unable to shrink filesystem at {partdev}. '
                     f'Unsupported filesystem "{fstype}".')


def ogExtendFs(disk, part):
    """
    Grow filesystem of a partition. Supports ext4 and ntfs partitions.
    Unsupported filesystem or invalid paths don't raise an exception,
    instead this method logs a warning message and does nothing.
    """
    partdev = get_partition_device(disk, part)
    fstype = get_filesystem_type(partdev)

    umount(partdev)
    if fstype == 'ext4':
        _extend_resize2fs(partdev)
    elif fstype == 'ntfs':
        _extend_ntfsresize(partdev)
    else:
        logging.warn(f'Unable to grow filesystem at {partdev}. '
                     f'Unsupported filesystem "{fstype}".')


def mkfs(fs, disk, partition, label=None):
    """
    Install any supported filesystem. Target partition is specified a disk
    number and partition number. This function uses utility functions to
    translate disk and partition number into a partition device path.

    If filesystem and partition are correct, calls the corresponding mkfs_*
    function with the partition device path. If not, ValueError is raised.
    """
    logging.debug(f'mkfs({fs}, {disk}, {partition}, {label})')
    fsdict = {
        'ext4': mkfs_ext4,
        'ntfs': mkfs_ntfs,
        'fat32': mkfs_fat32,
    }

    if fs not in fsdict:
        logging.warn(f'mkfs aborted, invalid target filesystem.')
        raise ValueError('Invalid target filesystem')

    try:
        partdev = get_partition_device(disk, partition)
    except ValueError as e:
        logging.warn(f'mkfs aborted, invalid partition.')
        raise e

    fsdict[fs](partdev, label)


def mkfs_ext4(partdev, label=None):
    if label:
        cmd = shlex.split(f'mkfs.ext4 -L {label} -F {partdev}')
    else:
        cmd = shlex.split(f'mkfs.ext4 -F {partdev}')
    with open('/tmp/command.log', 'wb', 0) as logfile:
        subprocess.run(cmd,
                       stdout=logfile,
                       stderr=STDOUT)


def mkfs_ntfs(partdev, label=None):
    if label:
        cmd = shlex.split(f'mkfs.ntfs -f -L {label} {partdev}')
    else:
        cmd = shlex.split(f'mkfs.ntfs -f {partdev}')
    with open('/tmp/command.log', 'wb', 0) as logfile:
        subprocess.run(cmd,
                       stdout=logfile,
                       stderr=STDOUT)


def mkfs_fat32(partdev, label=None):
    if label:
        cmd = shlex.split(f'mkfs.vfat -n {label} -F32 {partdev}')
    else:
        cmd = shlex.split(f'mkfs.vfat -F32 {partdev}')
    with open('/tmp/command.log', 'wb', 0) as logfile:
        subprocess.run(cmd,
                       stdout=logfile,
                       stderr=STDOUT)


def get_filesystem_type(partdev):
    """
    Get filesystem type from a partition device path.
    Raises RuntimeError when blkid exits with non-zero return code.
    """
    cmd = shlex.split(f'blkid -o value -s TYPE {partdev}')
    proc = subprocess.run(cmd, stdout=PIPE, encoding='utf-8')
    if proc.returncode != 0:
        raise RuntimeError(f'Error getting filesystem from {partdev}')
    return proc.stdout.strip()


def _reduce_resize2fs(partdev):
    cmd = shlex.split(f'resize2fs -fpM {partdev}')
    with open('/tmp/command.log', 'ab', 0) as logfile:
        subprocess.run(cmd, stdout=logfile, stderr=STDOUT)


def _reduce_ntfsresize(partdev):
    cmd_info = shlex.split(f'ntfsresize -Pfi {partdev}')
    proc_info = subprocess.run(cmd_info, stdout=subprocess.PIPE, encoding='utf-8')
    out_info = proc_info.stdout.strip()

    # Process ntfsresize output directly.
    # The first split operation leaves the wanted data at the second element of
    # the split ([1]). Finally do a second split with ' ' to get the data but
    # nothing else following it.
    #
    # In addition, increase by 10%+1K the reported shrink location on which the
    # filesystem can be resized according to ntfsresize.
    size = int(out_info.split('device size: ')[1].split(' ')[0])
    new_size = int(int(out_info.split('resize at ')[1].split(' ')[0])*1.1+1024)

    # Dry-run loop to test if resizing is actually possible. This is required by ntfsresize.
    returncode = 1
    while new_size < size and returncode != 0:
        cmd_resize_dryrun = shlex.split(f'ntfsresize -Pfns {new_size:.0f} {partdev}')
        proc_resize_dryrun = subprocess.run(cmd_resize_dryrun, stdout=subprocess.PIPE, encoding='utf-8')
        returncode = proc_resize_dryrun.returncode
        out_resize_dryrun = proc_resize_dryrun.stdout.strip()
        if 'Nothing to do: NTFS volume size is already OK.' in out_resize_dryrun:
            logging.warn('ntfsresize reports nothing to do. Is the target filesystem already shrunken?')
            break
        extra_size = int(out_resize_dryrun.split('Needed relocations : ')[1].split(' ')[0])*1.1+1024
        new_size += int(extra_size)

    if new_size < size:
        cmd_resize = shlex.split(f'ntfsresize -fs {new_size:.0f} {partdev}')
        with open('/tmp/command.log', 'ab', 0) as logfile:
            subprocess.run(cmd_resize, input='y', stderr=STDOUT, encoding='utf-8')


def _extend_resize2fs(partdev):
    cmd = shlex.split(f'resize2fs -f {partdev}')
    proc = subprocess.run(cmd)
    if proc.returncode != 0:
        raise RuntimeError(f'Error growing ext4 filesystem at {partdev}')


def _extend_ntfsresize(partdev):
    cmd = shlex.split(f'ntfsresize -f {partdev}')
    proc = subprocess.run(cmd, input=b'y')
    if proc.returncode != 0:
        raise RuntimeError(f'Error growing ntfs filesystem at {partdev}')