#
# 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 src.log import OgError

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)

    return True


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)
    if fstype == "unknown":
        return -1

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

    return ret

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)
    if fstype == "unknown":
        return -1

    umount(partdev)
    if fstype == 'ext4':
        _extend_resize2fs(partdev)
    elif fstype == 'ntfs':
        _extend_ntfsresize(partdev)
    else:
        logging.error(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,
        'linux-swap': mkfs_swap,
    }

    if fs not in fsdict:
        raise OgError(f'mkfs failed, unsupported target filesystem {fs}')

    try:
        partdev = get_partition_device(disk, partition)
    except ValueError as e:
        raise OgError(f'mkfs aborted: {e}') from 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 mkfs_swap(partdev, label=None):
    if label:
        cmd = shlex.split(f'mkswap -f -L {label} {partdev}')
    else:
        cmd = shlex.split(f'mkswap -f {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:
        log.error(f'Error getting filesystem from {partdev}')
        return "unknown"

    return proc.stdout.strip()


def _reduce_resize2fs(partdev):
    ret = -1
    cmd = shlex.split(f'resize2fs -fpM {partdev}')
    with open('/tmp/command.log', 'ab', 0) as logfile:
        proc = subprocess.run(cmd, stdout=logfile, stderr=STDOUT)
        if proc.returncode != 0:
            logging.error('failed to resize {partdev} with ext4')
            return -1

    return 0


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()

    if out_info.find('ERROR: NTFS is inconsistent. Run chkdsk') != -1:
        logging.error('NTFS is inconsistent. Run chkdsk /f on Windows then reboot TWICE!')
        return -1

    if proc_info.returncode != 0:
        logging.error(f'nfsresize {partdev} has failed with return code {proc_info.returncode}')
        return -1

    # 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.
    def parse_ntfsresize(output_data, pattern):
        data_split = output_data.split(pattern)
        # If we fail to match pattern in the split then data_split will contain [output_data]
        if len(data_split) == 1:
            raise OgError(f'nfsresize: failed to find: {pattern}')
        value_str = data_split[1].split(' ')[0]

        if not value_str.isdigit() or value_str.startswith('-'):
            raise OgError(f'nfsresize: failed to parse numeric value at {pattern}')
        return int(value_str)

    try:
        size = parse_ntfsresize(out_info, 'device size: ')
        new_size = parse_ntfsresize(out_info, 'resize at ')
        # Increase by 10%+1K the indicated reduction by which the file system
        # can be resized according to ntfsresize.
        new_size = int(new_size * 1.1 + 1024)
    except ValueError:
        return -1


    # Run ntfsresize with -n to to probe for the smallest size, this loop is
    # intentional. Acumulate size until ntfsresize in dry-run mode fails, then
    # use such size.
    while new_size < size:
        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')

        # valid new size found, stop probing
        if proc_resize_dryrun.returncode == 0:
            break

        out_resize_dryrun = proc_resize_dryrun.stdout.strip()
        if 'Nothing to do: NTFS volume size is already OK.' in out_resize_dryrun:
            logging.info('ntfsresize reports nothing to do. Is the target filesystem already shrunken?')
            break

        if out_resize_dryrun.find('Needed relocations : ') == -1:
            break

        try:
            extra_size = parse_ntfsresize(out_resize_dryrun, 'Needed relocations : ')
            # Add size padding
            extra_size = int(extra_size * 1.1 + 1024)
            new_size += extra_size
        except ValueError:
                return -1

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

    return 0

def _extend_resize2fs(partdev):
    cmd = shlex.split(f'resize2fs -f {partdev}')
    proc = subprocess.run(cmd)
    if proc.returncode != 0:
        raise OgError(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 OgError(f'Error growing ntfs filesystem at {partdev}')