#! /usr/bin/python3

import os, sys, argparse, glob, time, re
from subprocess import Popen, PIPE, call
from PyQt5.QtCore import QObject, QTranslator, QLocale
from PyQt5.QtWidgets import QApplication

class StickMaker(QObject):
    def __init__(self):
        QObject.__init__(self)
        self.startTime = time.time()
        self.pathSentence="PATH=/usr/sbin:/sbin:/usr/bin:/bin; "

    def hms(self):
        """
        returns the time elapsed so far, in format HH:MM:SS
        """
        t=int(time.time()-self.startTime)
        th=t//3600
        tmin=(t-3600*th)//60
        tsec=(t-3600*th-60*tmin)
        return f"{th:02d}:{tmin:02d}:{tsec:02d}"

    def flush(self, stderr=True, stdout=True):
        """
        flush standard outputs
        @param stderr don't flush sys.stderr if it's False (True by default)
        @param stdout don't flush sys.stdout if it's False (True by default)
        """
        if stderr: sys.stderr.flush()
        if stdout: sys.stdout.flush()
        return

    def log(self, s, timestamp=True):
        """
        print some information for stdout
        @param s a string to log
        @param timestamp True (by default) implies that there will be
        a timestamp
        """
        if timestamp:
            print(self.tr("[At {0}]    {1}").format(self.hms(),s))
        else:
            print(s)
        sys.stdout.flush()
        return

    def run(self):
        """
        The main method
        """
        if os.geteuid() != 0:
            self.log(
                self.tr("You must be root to launch {}").format(sys.argv[0]),
                timestamp=False
            )
            sys.exit(1)


        parser = argparse.ArgumentParser(description=self.tr('clone a Debian-Live distribution'))
        parser.add_argument('-f','--vfat-size', help=self.tr("size of the VFAT partition, in GB"))
        parser.add_argument('-n','--no-persistence-seed', help=self.tr("do'nt use the file rw.tar.gz to seed the persistence"), action='store_true')
        parser.add_argument('-s','--source', help=self.tr('source of the distribution (required); the keyword OWN means that we are cloning from a Debian-Live disk'), metavar="source", required=True)
        parser.add_argument('device', help=self.tr('the targetted USB stick'))
        args = parser.parse_args()

        ################### unmount targetted partitions if any #############

        cmd=f"mount | grep {args.device} | awk '{{print $1}}'"
        p=Popen(self.pathSentence+cmd, shell=True, stdout=PIPE, stderr=PIPE)
        mountpoints,_ = p.communicate()
        mounts=mountpoints.decode("utf-8").split()
        if mounts:
            self.log(self.tr("Unmounting {}").format(', '.join(mounts)))
        for m in mounts:
            call(self.pathSentence+f"umount {m}", shell=True)
        self.flush()

        ################### check the sizes #################################

        cmd = "lsblk| grep {}| grep disk| awk '{{print $4}}'".format(args.device.replace("/dev/",""))
        p = Popen(cmd, shell=True, stdout=PIPE)
        size, _ = p.communicate()
        size = size.decode("utf-8").strip().replace(",",".")
        assert(re.match(r"\d+(\.\d+)?G", size))
        size=float(size.replace("G",""))
        iso_size = os.stat(args.source).st_size//1024**3
        min_persistence = 1 # at least 1 GB for the persistence

        if iso_size + min_persistence > size:
            self.log(self.tr("The disk is too small ({} GiB) to contain the ISO image ({} GiB) plus at least 1 GB for the persistence").format(size, iso_size))
            assert (size > iso_size + min_persistence) # to stop the process

        if iso_size + float(args.vfat_size) + min_persistence > size:
            new_vfat_size = size - iso_size - min_persistence
            self.log(self.tr("The size of the VFAT partition has been shrinked down to {} GB").format(new_vfat_size))
            args.vfat_size = str(new_vfat_size)
        

        ################### copy the ISO-hybrid image #######################

        self.log(self.tr("Copying {0} to {1} ...").format(args.source, args.device))

        if args.source.upper()=="OWN":
            mountpoint = glob.glob("/usr/lib/live/mount/persistence/sd*1")[0]
            sourceDevice=mountpoint.replace("/usr/lib/live/mount/persistence","/dev")[:-1]
            ### when copying from own, the count of 4M blocks is computed
            ### from the last sector used by /dev/sd.1
            cmd=f"sfdisk -l {sourceDevice}| grep /dev/sd.1| awk '{{print $4}}'"
            p=Popen(self.pathSentence+cmd, shell=True, stdout=PIPE, stderr=PIPE)
            sectors,_ = p.communicate()
            sectors=int(sectors.decode("utf-8").strip())
            count = 1+sectors//8192 # convert 512 B sectors to 4 MiB blocks
            
            cmd = f"dd if={sourceDevice} of={args.device} status=progress bs=4M oflag=dsync count={count}"
        else:
            ### Then, the boot sector bears information for two partitions.
            cmd=f"dd if={args.source} of={args.device} status=progress bs=4M oflag=dsync"

        call(self.pathSentence+cmd, shell=True)
        self.flush()
        
        cmd=f"sync; sleep 4; partprobe {args.device}"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        inputs="" ### command lines for fdisk

        if args.source.upper()=="OWN":
            # if there are partitions #4 and #3 in the table, delete them
            for num in (4, 3):
                cmd=f"sfdisk -l {args.device}| grep {args.device}{num}"
                p=Popen(self.pathSentence+cmd, shell=True, stdout=PIPE, stderr=PIPE)
                reply, _ = p.communicate()
                if reply.decode("utf-8").strip(): # the partition #{num} does exist
                    self.log(self.tr("Delete the partition {}{}, to create it again").format(args.device, num))
                    inputs+=f"""\
d
{num}
"""

        ############# adds partition #3 for the VFAT data   ##########
        ############# adds partition #4 for the persistence ##########
        self.log(self.tr("Create the partition {}3 for VFAT data, using {}GiB").format(args.device, args.vfat_size))
        self.log(self.tr("Create the partition {}4 to support persistence").format(args.device))
        self.log(self.tr("This partition is adjusted to use all the remaining space on {}").format(args.device))
        inputs+="""\
n
p
3

+{}G
t
3
0b
n
p
4



w
""".format(args.vfat_size)
        p=Popen(self.pathSentence+f"fdisk {args.device}", shell=True, stdout=PIPE, stderr=PIPE, stdin=PIPE)
        p.stdin.write(inputs.encode("utf-8"))

        message,err=p.communicate()
        self.log(message.decode("utf-8"))
        self.log(err.decode("utf-8"))

        cmd=f"sleep 6; partprobe {args.device}"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        ################ format VFAT   ####################################

        cmd = f"mkfs.vfat -n DATA {args.device}3"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        ################ format (ext4) ####################################

        cmd = f"mkfs.ext4 -L persistence -F {args.device}4"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        ################## mount the 4th partition ########################

        cmd = f"mkdir -p /tmp/{args.device}4; mount {args.device}4 /tmp/{args.device}4"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        ################# write the file persistence.conf ################

        self.log(self.tr("Create the file persistence.conf on the fourth partition"))
        with open(f"/tmp/{args.device}4/persistence.conf", "w") as outfile:
            outfile.write("/ union\n")

        ################## eventually seed the persistence partition ######

        if args.no_persistence_seed==False and os.path.exists("rw.tar.gz"):
            self.log(self.tr("Pre-seed the fourth partition with rw.tar.gz"))
            cmd = f"zcat rw.tar.gz | (cd /tmp/{args.device}4/; tar xf -)"
            call(self.pathSentence+cmd, shell=True)
        self.flush()

        #################### little cleanup ###############################

        cmd = f"umount /tmp/{args.device}4; rmdir /tmp/{args.device}4; partprobe {args.device}"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        cmd = f"sfdisk -l {args.device}"
        call(self.pathSentence+cmd, shell=True)
        self.flush()

        self.log(
            "\n=======================================================\n",
            timestamp=False
        )
        self.log(
            self.tr("Ready. Total time: {}").format(self.hms()),
            timestamp=False
        )

if __name__=="__main__":
    app = QApplication(sys.argv)
    # i18n stuff
    locale = QLocale.system().name()
    translation="live_clone_{}.ts".format(locale)
    langPath=os.path.join(os.path.abspath(os.path.dirname(__file__)),"lang",translation)
    translator = QTranslator(app)
    translator.load(langPath)
    app.installTranslator(translator)
    #############
    maker = StickMaker()
    maker.run()
