#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  auto_partition.py
#
#  This file was forked from Cnchi (graphical installer from Antergos)
#  Check it at https://github.com/antergos
#
#  Copyright 2013 Antergos (http://antergos.com/)
#  Copyright 2013 Manjaro (http://manjaro.org)
#
#  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; either version 2 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 General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.

import os
import subprocess
import logging
import show_message as show
import parted3.partition_module as pm
import parted3.fs_module as fs
import parted3.lvm as lvm
import parted3.used_space as used_space

""" AutoPartition class """

# Partition sizes are in MiB
MAX_ROOT_SIZE = 30000

# TODO: This higly depends on the selected DE! Must be taken into account.
# KDE needs 4.5 GB for its files. Need to leave extra space also.
MIN_ROOT_SIZE = 6500

def check_output(command):
    """ Calls subprocess.check_output, decodes its exit and removes trailing \n """
    return subprocess.check_output(command.split()).decode().strip("\n")


def printk(enable):
    """ Enables / disables printing kernel messages to console """
    with open("/proc/sys/kernel/printk", "w") as fpk:
        if enable:
            fpk.write("4")
        else:
            fpk.write("0")

def unmount_all(dest_dir):
    """ Unmounts all devices that are mounted inside dest_dir """
    swaps = subprocess.check_output(["swapon", "--show=NAME", "--noheadings"]).de
code().split("\n")
    for name in filter(None, swaps):
        if "/dev/zram" not in name:
            subprocess.check_call(["swapoff", name])

    mount_result = subprocess.check_output("mount").decode().split("\n")

    # Umount all devices mounted inside dest_dir (if any)
    dirs = []
    for mount in mount_result:
        if dest_dir in mount:
            directory = mount.split()[0]
            # Do not unmount dest_dir
 now (we will do it later)
            if directory is not dest_dir:
                dirs.append(directory)

    for directory in dirs:
        logging.warning(_("Unmounting %s"), directory)
        try:
            subprocess.call(["umount", directory])
        except Exception:
            logging.warning(_("Unmounting %s failed. Trying lazy arg."), directory)
            subprocess.call(["umount", "-l", directory])

    # Now is the time to unmount the device that is mounted in dest_dir (if any)
    if dest_dir in mount_result:
        logging.warning(_("Unmounting %s"), dest_dir)
        try:
            subprocess.call(["umount", dest_dir])
        except Exception:
            logging.warning(_("Unmounting %s failed. Trying lazy arg."), dest_dir)
            subprocess.call(["umount", "-l", dest_dir])

    # Remove all previous LVM volumes
    # (it may have been left created due to a previous failed installation)
    try:
        lvolumes = check_output("lvs -o lv_name,vg_name --noheading").split("\n")
        if len(lvolumes[0]) > 0:
            for lvolume in lvolumes:
                if len(lvolume) > 0:
                    (lvolume, vgroup) = lvolume.split()
                    lvdev = "/dev/" + vgroup + "/" + lvolume
                    subprocess.check_call(["wipefs", "-af", lvdev])
                    subprocess.check_call(["lvremove", "-f", lvdev])

        vgnames = check_output("vgs -o vg_name --noheading").split("\n")
        if len(vgnames[0]) > 0:
            for vgname in vgnames:
                vgname = vgname.strip()
                if len(vgname) > 0:
                    subprocess.check_call(["vgremove", "-f", vgname])

        pvolumes = check_output("pvs -o pv_name --noheading").split("\n")
        if len(pvolumes[0]) > 0:
            for pvolume in pvolumes:
                pvolume = pvolume.strip(" ")
                subprocess.check_call(["pvremove", "-f", pvolume])

    except subprocess.CalledProcessError as err:
        logging.warning(_("Can't delete existent LVM volumes (see below)"))
        logging.warning(err)

    # Close LUKS devices (they may have been left open because of a previous failed installation)
    try:
        if os.path.exists("/dev/mapper/cryptManjaro"):
            subprocess.check_call(["cryptsetup", "luksClose", "/dev/mapper/cryptManjaro"])
        if os.path.exists("/dev/mapper/cryptManjaroHome"):
            subprocess.check_call(["cryptsetup", "luksClose", "/dev/mapper/cryptManjaroHome"])
    except subprocess.CalledProcessError as err:
        logging.warning(_("Can't close LUKS devices (see below)"))
        logging.warning(err)

class AutoPartition(object):
    """ Class used by the automatic installation method """
    def __init__(self, dest_dir, auto_device, use_luks, use_lvm, luks_key_pass, use_home, callback_queue):
        """ Class initialization """
        self.dest_dir = dest_dir
        self.auto_device = auto_device
        self.luks_key_pass = luks_key_pass
        # Use LUKS encryption
        self.luks = use_luks
        # Use LVM
        self.lvm = use_lvm
        # Make home a different partition or if using LVM, a different volume
        self.home = use_home

        # Will use these queue to show progress info to the user
        self.callback_queue = callback_queue

        self.efi = False
        if os.path.exists("/sys/firmware/efi"):
            self.efi = True

        self.separate_boot = use_luks or use_lvm or self.efi
        logging.debug( "luks is " + str(use_luks) + ", lvm is " + str(use_lvm) \
                       + " and efi is " + str(self.efi) \
                       + "\ntherefore separate_boot is " + str(self.separate_boot))

    def mkfs(self, device, fs_type, mount_point, label_name, fs_options="", btrfs_devices=""):
        """ We have two main cases: "swap" and everything else. """
        logging.debug("Will mkfs " + device + " as " + fs_type)
        if fs_type == "swap":
            try:
                swap_devices = check_output("swapon -s")
                if device in swap_devices:
                    subprocess.check_call(["swapoff", device])
                subprocess.check_call(["mkswap", "-L", label_name, device])
                subprocess.check_call(["swapon", device])
            except subprocess.CalledProcessError as err:
                logging.warning(err.output)
        else:
            mkfs = {"xfs": "mkfs.xfs %s -L %s -f %s" % (fs_options, label_name, device),
                    "jfs": "yes | mkfs.jfs %s -L %s %s" % (fs_options, label_name, device),
                    "reiserfs": "yes | mkreiserfs %s -l %s %s" % (fs_options, label_name, device),
                    "ext2": "mkfs.ext2 -q  %s -F -L %s %s" % (fs_options, label_name, device),
                    "ext3": "mkfs.ext3 -q  %s -F -L %s %s" % (fs_options, label_name, device),
                    "ext4": "mkfs.ext4 -q  %s -F -L %s %s" % (fs_options, label_name, device),
                    "btrfs": "mkfs.btrfs %s -L %s %s" % (fs_options, label_name, btrfs_devices),
                    "nilfs2": "mkfs.nilfs2 %s -L %s %s" % (fs_options, label_name, device),
                    "ntfs-3g": "mkfs.ntfs %s -L %s %s" % (fs_options, label_name, device),
                    "vfat": "mkfs.vfat %s -n %s %s" % (fs_options, label_name, device)}

            # Make sure the fs type is one we can handle
            if fs_type not in mkfs.keys():
                txt = _("Unknown filesystem type %s") % fs_type
                logging.error(txt)
                show.error(txt)
                return

            command = mkfs[fs_type]

            try:
                subprocess.check_call(command.split())
            except sub