#
# Copyright (C) 2020-2024 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 hashlib
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, readonly=False):
    """
    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):
        try:
            os.mkdir(target)
        except OSError as e:
            logging.error(f'mkdir operation failed. Reported {e}')
            return False

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

    return True


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

    Return true if exit code is 0. False otherwise.
    """
    if readonly:
        cmd = f'mount -o ro {source} {target}'
    else:
        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 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)
    elif fstype == 'vfat':
        ret = 0
    else:
        ret = -1
        logging.error(f'Unable to shrink filesystem at {partdev}. '
		      f'Unsupported filesystem "{fstype}".')

    return ret

def extend_filesystem(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)
    elif fstype == 'vfat':
        pass
    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

    return fsdict[fs](partdev, label)


def mkfs_ext4(partdev, label=None):
    err = -1
    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:
        ret = subprocess.run(cmd,
                             stdout=logfile,
                             stderr=STDOUT)
        err = ret.returncode
        if ret.returncode != 0:
            logging.error(f'mkfs.ext4 reports return code {ret.returncode} for {partdev}')

    return err


def mkfs_ntfs(partdev, label=None):
    err = -1
    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:
        ret = subprocess.run(cmd,
                             stdout=logfile,
                             stderr=STDOUT)
        err = ret.returncode
        if ret.returncode != 0:
            logging.error(f'mkfs.ntfs reports return code {ret.returncode} for {partdev}')

    return err


def mkfs_fat32(partdev, label=None):
    err = -1
    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:
        ret = subprocess.run(cmd,
                             stdout=logfile,
                             stderr=STDOUT)
        err = ret.returncode
        if ret.returncode != 0:
            logging.error(f'mkfs.vfat reports return code {ret.returncode} for {partdev}')

    return err


def mkfs_swap(partdev, label=None):
    err = -1
    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:
        ret = subprocess.run(cmd,
                             stdout=logfile,
                             stderr=STDOUT)
        err = ret.returncode
        if ret.returncode != 0:
            logging.error(f'mkswap reports return code {ret.returncode} for {partdev}')

    return err


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:
        logging.error(f'Error getting filesystem from {partdev}')
        return "unknown"

    return proc.stdout.strip()


def _reduce_resize2fs(partdev):
    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(f'Failed to resize ext4 filesystem in {partdev}')
            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}')

def compute_md5(path, bs=2**20):
    if not os.path.exists(path):
        raise OgError(f"Failed to calculate checksum, image file {path} does not exist")

    m = hashlib.md5()
    try:
        with open(path, 'rb') as f:
            while True:
                buf = f.read(bs)
                if not buf:
                    break
                m.update(buf)
    except Exception as e:
        raise OgError(f'Failed to calculate checksum for {path}: {e}') from e
    return m.hexdigest()