#!/usr/bin/python3
#
#  imported from project Mintstick, by Clement Lefebvre
#  https://github.com/linuxmint/mintstick/
#
#  Copyright (C) 2013 THE isodumper'S COPYRIGHT HOLDER
#  This file is distributed under the same license as the isodumper package.
#
#  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 St, Fifth Floor, Boston, MA 02110-1301, USA.
#
#  Author isodumper software= papoteur <papoteur@mageialinux-online.org>

import sys
import getopt

from subprocess import call

sys.path.append("/usr/lib/isodumper")


def raw_format(device_path, fstype, volume_label, uid, gid):
    import parted

    # First erase MBR and partition table , if any
    call(["dd", "if=/dev/zero", "of=%s" % device_path, "bs=512", "count=1"])

    device = parted.getDevice(device_path)

    # Create a default partition set up
    disk = parted.freshDisk(device, "msdos")
    try:
        disk.commit()
    except:
        # Unable to write the new partition
        print("Can't create partition. The device could be read-only.")
        sys.exit(5)
    regions = disk.getFreeSpaceRegions()

    if len(regions) > 0:
        # print "Build partition"
        # Define size
        start = parted.sizeToSectors(1, "MiB", device.sectorSize)
        # print "start %s" % start
        end = device.getLength() - start - 1024
        # print end

        # Alignment
        # cylinder = device.endSectorToCylinder(end)
        # end = device.endCylinderToSector(cylinder)
        # print end
        try:
            geometry = parted.Geometry(device=device, start=start, end=end)
        except:
            print("Geometry error - Can't create partition")
            sys.exit(5)

        # fstype
        _fstype = fstype
        if _fstype == "exfat":
            _fstype = "ntfs"
        fs = parted.FileSystem(type=_fstype, geometry=geometry)

        # Create partition
        partition = parted.Partition(
            disk=disk, type=parted.PARTITION_NORMAL, geometry=geometry, fs=fs
        )
        constraint = parted.Constraint(exactGeom=geometry)
        disk.addPartition(partition=partition, constraint=constraint)
        disk.commit()

        # Format partition according to the fstype specified
        if fstype == "fat32":
            call(["mkdosfs", "-F", "32", "-n", volume_label, partition.path])
        elif fstype == "exfat":
            call(["mkfs.exfat", "-n", volume_label, partition.path])
        elif fstype == "ntfs":
            call(["mkntfs", "-f", "-L", volume_label, partition.path])
        elif fstype == "ext4":
            call(
                [
                    "mkfs.ext4",
                    "-E",
                    "root_owner=%s:%s" % (uid, gid),
                    "-L",
                    volume_label,
                    partition.path,
                ]
            )
    sys.exit(0)


def main():
    # parse command line options
    try:
        opts, args = getopt.getopt(
            sys.argv[1:],
            "hd:f:l:u:g:",
            ["help", "device=", "filesystem=", "label=", "uid=", "gid="],
        )
    except getopt.error as msg:
        print(msg)
        print("for help use --help")
        sys.exit(2)

    for o, a in opts:
        if o in ("-h", "--help"):
            print("Usage: %s -d device -f filesystem -l volume_label\n" % sys.argv[0])
            print("-d|--device          : device path")
            print("-f|--filesystem      : filesystem\n")
            print("-l|--label           : volume label\n")
            print("-u|--uid             : uid of user\n")
            print("-g|--gid             : gid of user\n")
            print(
                'Example : %s -d /dev/sdj -f fat32 -l "USB Stick" -u 1000 -g 1000'
                % sys.argv[0]
            )
            sys.exit(0)
        elif o in ("-d"):
            device = a
        elif o in ("-f"):
            if a not in ["fat32", "exfat", "ntfs", "ext4"]:
                print("Specify fat32, exfat, ntfs or ext4")
                sys.exit(3)
            fstype = a
        elif o in ("-l"):
            label = a
        elif o in ("-u"):
            uid = a
        elif o in ("-g"):
            gid = a

    argc = len(sys.argv)
    if argc < 11:
        print("Too few arguments")
        print("for help use --help")
        exit(2)

    raw_format(device, fstype, label, uid, gid)


if __name__ == "__main__":
    main()