"""
Defines wizards related to actions of live-clone
"""
from PyQt5.QtWidgets import QWizard, QWizardPage, QLineEdit, \
    QHBoxLayout, QVBoxLayout, \
    QRadioButton, QPlainTextEdit, QTreeView, QWidget, QCheckBox, \
    QMessageBox, QPushButton, QFileDialog, QLabel, QFrame
from PyQt5.QtGui import QStandardItemModel, QStandardItem
from PyQt5.QtCore import Qt

from tools import MountPoint, FileTreeSelectorModel
from monitor import Monitor, CloneMonitor

from subprocess import Popen, PIPE, call
from datetime import datetime
import os, re, glob
from simul import loopDict

def createVfatExt4Partitions(device, size="4G"):
    """
    create a VFAT and a EXT4 partition on a device.
    first find the number of the first VFAT partition, and erase any
    partition with this number or a greater number.
    Then create a VFAT partiton with size `size`, and a EXT4 partition
    to use all the space remaining on the device. Finally format both
    partitons according to their types
    @param device a string, like "/dev/sdc"
    @param size a string, defaults to "4G"
    @return the number of the last created partition (i.e. the EXT4 partition's
    number)
    """
    cmd="fdisk -l {}".format(device)
    p=Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
    out, err= p.communicate()
    out=out.decode("utf-8")
    FATlines=[]
    for l in out.split("\n"):
        m=re.match(r"^(/dev/\S+)\s.*FAT",l)
        if m:
            # remember the FAT partition and unmount it
            FATlines.append(m.group(1))
            call(f"umount {m.group(1)}", shell=True)
        else:
            m=re.match(r"^(/dev/\S+)\s.*",l)
            if m:
                # unmount other partitions too
                call(f"umount {m.group(1)}", shell=True)
    assert (FATlines) # there *must* be a first VFAT partition
    m=re.match(r"/dev/\D+(\d+)",FATlines[0])
    fatNum=int(m.group(1))
    inputs=f"""\
d

n
p
{fatNum}

+4G
t
0c
n
p
{fatNum+1}



w

"""
    p=Popen(f"fdisk {device}", shell=True, stdout=PIPE, stderr=PIPE, stdin=PIPE)
    p.stdin.write(inputs.encode("utf-8"))
    message,err=p.communicate()
    call(f"sleep 6; partprobe {device}", shell=True)
    persistenceNum = fatNum+1
    call(f"mkfs.vfat -n DATA {device}{fatNum}", shell=True)
    call(f"mkfs.ext4 -L persistence -F {device}{persistenceNum}", shell=True)
    return persistenceNum

def new_packages(d, persistenceNum):
    """
    Work out a list of newer packages installed in the persistence
    zone of a debian-live Disk
    @param d a disk given as a /dev/sdx string
    @param persistenceNum the number of the persistence partition
    @return a text with on package by line
    """
    with MountPoint(d+"1") as base, MountPoint(d+str(persistenceNum)) as pers:
        cmd="""find "{pers}/rw/var/lib/dpkg/info" -name "*.list" -newer "{base}/live/filesystem.packages" | sed -e "s%{pers}/rw/var/lib/dpkg/info/%%" -e 's/.list$//' -e 's/:.*//'""".format(base=base, pers=pers)
        p=Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
        out, err= p.communicate()
    return out.decode("utf-8")

class ToolWizard(QWizard):
    """
    The wizard associated to the "Tools" button of live-clone
    """

    def __init__(self, mw):
        """
        the constructor
        @param mw a MyMain (<= QMainWindow) instance
        """
        QWizard.__init__(self, mw)
        self.setWindowTitle(self.tr("Repair and/or manage a Freeduc USB stick"))
        page1=QWizardPage(self)
        page1.setTitle(self.tr("Select one disk"))
        page1.setSubTitle(self.tr("Which disk do you want to repair or to save?"))
        target=QLineEdit(page1)
        target.hide()
        page1.registerField("target*", target)
        layout=QVBoxLayout()
        radios={}
        def setTarget(checked):
            for r in radios:
                if r.isChecked():
                    page1.setField("target", r.text())
            return
        for d in mw.not_self_keys(format="udev"):
            rb=QRadioButton(d, parent=page1)
            rb.toggled.connect(setTarget)
            layout.addWidget(rb)
            radios[rb]=d
        page1.setLayout(layout)
        self.addPage(page1)
        ## assert page1.field("target")
        page2=QWizardPage(self)
        disk=self.field("target")
        persistenceNum=4 # default number for the persistence partition
        page2.setTitle(self.tr("List of additional packages in {}").format(disk))
        page2.setSubTitle(self.tr("Here is the list of packages which you can restore later. You can edit the text before saving it."))
        layout=QVBoxLayout()
        ed=QPlainTextEdit(page2)
        layout.addWidget(ed)
        newpackages=True
        def page2init(foo=None):
            nonlocal newpackages
            nonlocal persistenceNum
            disk=self.field("target")
            ## find the persistence partition if any
            diskname=disk.replace("/dev/","")
            cmd="ls -l /dev/disk/by-label| grep persistence | grep "+diskname
            p=Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
            out, err= p.communicate()
            if out:
                persistenceNum=out.decode("utf8").strip()[-1]
            else:
                newpackages=False
                return # no persistence, no need to develop the wizard further

            # we know the number of the persistence partition
            newpackages=new_packages(disk, persistenceNum)
            if not newpackages:
                packages=self.tr("There are no new packages installed in the persistence partition")
            else:
                packages=self.tr("""\
ADDITIONAL PACKAGES
===================
""") + "{}\n{}\n".format(disk,len(disk)*"-") + newpackages
            c=ed.textCursor()
            c.insertText(packages)
            ed.moveCursor(c.Start)
            ed.ensureCursorVisible()
            return
        page2.setLayout(layout)
        page2.initializePage=page2init
        def page2Validate(foo=None):
            if not newpackages:
                return True
            proposed_fname = "packages-{}.md".format(
                datetime.now().strftime("%Y-%m-%d--%H-%M")
            )
            fname, _ = QFileDialog.getSaveFileName(
                self, self.tr("Save Packages File"),
                proposed_fname,
                self.tr("Markdow files (*.md);;Plain text files (*.txt);;Any file (*)")
            )
            if fname:
                with open(fname,"w") as outfile:
                    outfile.write(ed.toPlainText())
                if "SUDO_UID" in os.environ and "SUDO_GID" in os.environ:
                    ## it is possible to downgrade the owner of the file
                    os.chown(
                        fname,
                        int(os.environ["SUDO_UID"]),
                        int(os.environ["SUDO_GID"])
                    )
            return True
        page2.validatePage=page2Validate
        self.addPage(page2)
        page3=QWizardPage(self)
        page3.setTitle(self.tr("Select personal data to save"))
        page3.setSubTitle(self.tr('Please check the tree of data which you want to backup. If you check the topmost node "user", all and every personal data is chosen, including Mozilla cache and desktop preferences. You may prefer to check subtrees. The subtrees are taken in account only when no node is checked above.'))
        layout=QVBoxLayout()
        tree=QTreeView(page3)
        layout.addWidget(tree)
        mp=None
        model=None
        exists=False
        def page3Init(foo=None):
            nonlocal mp, model, exists
            mp=MountPoint(self.field("target")+str(persistenceNum))
            mp.enter_()
            if mp:
                rootpath=os.path.join(str(mp), "rw", "home")
                exists = os.path.exists(rootpath)
                if exists:
                    model = FileTreeSelectorModel(rootpath=rootpath)
                    model.setRootPath(rootpath)
                    tree.setModel(model)
                    tree.setRootIndex(model.parent_index)

                    tree.setAnimated(False)
                    tree.setIndentation(20)
                    tree.setSortingEnabled(True)
                else:
                    ## there are no personal data
                    model = QStandardItemModel()
                    model.setHorizontalHeaderLabels(
                        [self.tr('Nothing is available')])
                    tree.setModel(model)
                    tree.setUniformRowHeights(True)

                    node1 = QStandardItem(
                        self.tr('No personal data in the persistence partition'))
                    model.appendRow(node1)
            return
        def page3validate(foo=None):
            nonlocal mp, exists
            if not exists:
                return True
            to_save=[]
            def cb(i, stack):
                if model.checkState(i)==Qt.Checked:
                    pathdata=[index.data() for index in stack+[i]]
                    to_save.append(os.path.join(*pathdata))
                return
            i=model.parent_index
            model.traverseDirectoryWhileUnchecked(i,callback=cb)
            ### ask for a tgz file name and make a tarball into it
            proposed_fname = "DATA-{}.tgz".format(
                datetime.now().strftime("%Y-%m-%d--%H-%M")
            )
            fname=""
            if to_save:
                fname, _ = QFileDialog.getSaveFileName(
                    self, self.tr("Save personal data archive"),
                    proposed_fname,
                    self.tr("TAR-GZIP archives (*.tgz,*.tar.gz);;Any file (*)")
                )
            if fname:
                ### launch the command asynchronously with a Monitor
                command = "/usr/bin/tar"
                args=[
                    "czvf",
                    fname
                ] + [os.path.join(str(mp), "rw", p) for p in to_save]
                ### add a new tab and a monitor
                n=mw.ui.tabWidget.count()
                title=self.tr("Save personal persistence data to {}").format(fname)
                newWidget=QWidget()
                mw.ui.tabWidget.addTab(
                    QWidget(),
                    self.tr("Save to {}").format(fname)
                )
                rindex=mw.rindexFromDev(self.field("target"))
                if rindex==None:
                    return True
                # rindex is defined, it bears the row number for the device
                def makeFinishCallback():
                    return lambda: mp.exit_()
                m=Monitor(
                    command, args, mw, n, rindex,
                    finishCallback=makeFinishCallback(),
                    environ=os.environ,
                    createfiles=[fname],
                    actionMessage=self.tr("save-data"),
                    buttonMessage=self.tr("Close the log of the archive"),
                )
                mw.monitors[m]=self.field("target")
                m.start()
            return True
        page3.initializePage=page3Init
        page3.validatePage=page3validate
        page3.setLayout(layout)
        self.addPage(page3)
        page4=QWizardPage(self)
        page4.setTitle(self.tr("Select personal data to save"))
        page4.setSubTitle(self.tr('Now, you have carefully backuped sensistive data, you may erase the persistence; however, you can still cancel it.'))
        def page4validate(foo=None):
            if self.field("ok1") and  self.field("ok2"):
                pers=self.field("target")+str(persistenceNum)
                call ("umount {}".format(pers), shell=True)
                def page4CallBackMaker():
                    def callback():
                        nonlocal persistenceNum
                        must_format = False # do we have to format the device?
                        # first check whether the persistence partition exists
                        with MountPoint(self.field("target")+str(persistenceNum)) as mp:
                            if not mp.capacity:
                                must_format = True
                        # format the device if necessary
                        if must_format:
                            persistenceNum = createVfatExt4Partitions(self.field("target"))
                        # then create the persistence.conf file
                        with MountPoint(self.field("target")+str(persistenceNum)) as mp:
                            confFile=os.path.join(str(mp),"persistence.conf")
                            call("echo '/ union' > {}".format(confFile), shell=True)
                        QMessageBox.information(
                            self,
                            self.tr("Persistence erased"),
                            self.tr("Persistence data have been erased from {}").format(self.field("target")+str(persistenceNum))
                        )
                        return
                    return callback
                command="mkfs"
                args= [
                    "-F",
                    "-L",
                    "persistence",
                    "-t",
                    "ext4",
                    pers
                ]
                n=mw.ui.tabWidget.count()
                title=self.tr("Format the partition {}").format(pers)
                newWidget=QWidget()
                mw.ui.tabWidget.addTab(
                    QWidget(),
                    self.tr("Format the partition {}").format(pers)
                )
                rindex=mw.rindexFromDev(self.field("target"))
                if rindex==None:
                    return True                
                m=Monitor(command, args, mw, n, rindex,
                          actionMessage="Formating",
                          finishCallback=page4CallBackMaker()
                )
                mw.monitors[m]=self.field("target")
                m.start()
            return True
        layout=QVBoxLayout()
        request1CB=QCheckBox(self.tr("Yes, I am sure that I want to erase the persistence."))
        request2CB=QCheckBox(self.tr("Yes again, I am completely sure!"))
        page4.registerField("ok1", request1CB)
        page4.registerField("ok2", request2CB)
        layout.addWidget(request1CB)
        layout.addWidget(request2CB)
        page4.setLayout(layout)
        page4.validatePage=page4validate
        self.addPage(page4)
        return
    
class CloneWizard(QWizard):
    """
    The wizard associated to the "Clone" button of live-clone
    """

    def __init__(self, mw):
        """
        the constructor
        @param mw a MyMain (<= QMainWindow) instance
        """
        QWizard.__init__(self, mw)
        self.setWindowTitle(self.tr("Clone a Freeduc system to USB drives"))
        page1=QWizardPage(self)
        page1.setTitle(self.tr("Select the target disks"))
        page1.setSubTitle(self.tr("Please check the disks where you want to clone the system. Warning: all data exiting on those disks will be erased, you can still Escape from this process."))
        selectedDisks=QLineEdit()
        selectedDisks.setReadOnly(True)
        selectedDisks.hide() # not visible, it bears juste a calculated field
        page1.registerField("selectedDisks*", selectedDisks)
        boxes={}
        def updateSelectedDisks(checkState=-1):
            """
            callback for changes in checkboxes
            """
            selected=[d for b,d in boxes.items() \
                      if b.checkState()==Qt.Checked]
            sel=",".join(selected)
            page1.setField("selectedDisks", sel)
            if checkState==-1:
                #called before the wizard is initialized
                selectedDisks.setText(sel)
            return
        vfatLineEdit = QLineEdit("0")
        page1.registerField("vfat", vfatLineEdit)
        def updateVFATsize(size = "0"):
            """
            callback for changes VFAT size lineEdit widget
            """
            value = ""
            for char in size:
                try:
                    v = float(value+char)
                    value = value + char
                except:
                    break
            size=value
            if size == "":
                size = "0"
            page1.setField("vfat", size)
            vfatLineEdit.setText(size)
            return
        vfatLineEdit.textEdited.connect(updateVFATsize)
        layout=QVBoxLayout()
        for disk in mw.not_self_keys(format="udev"):
            shortParts=mw.lesCles.parts_summary(disk)
            chb=QCheckBox("{} [{}]".format(
                disk,
                ", ".join(shortParts))
            )
            boxes[chb]=disk
            chb.stateChanged.connect(updateSelectedDisks)
            layout.addWidget(chb)
        layout.addWidget(selectedDisks)
        layout.addStretch()
        fr = QFrame()
        frLayout = QHBoxLayout()
        frLayout.addWidget(QLabel(self.tr("Size of a VFAT partition (GB):")))
        frLayout.addWidget(vfatLineEdit)
        frLayout.addStretch()
        fr.setLayout(frLayout)
        layout.addWidget(fr)
        page1.setLayout(layout)
        self.addPage(page1)
        sourceEdit=QLineEdit()
        sourceEdit.setReadOnly(True)
        page2=QWizardPage(self)
        page2.registerField("source*", sourceEdit)
        ownClone=os.path.exists("/usr/lib/live/mount/persistence")
        def page2init(foo=None):
            if ownClone:
                sourceEdit.setText("AUTOCLONE (Freeduc disks can clone themselves natively)")
            return
        page2.initializePage=page2init
        page2.setTitle(self.tr("Source of the core system"))
        if ownClone:
            page2.setSubTitle(self.tr("The program is running from a Freeduc GNU-Linux system"))
        else:
            page2.setSubTitle(self.tr("The program is running from a plain GNU-Linux system"))
        layout=QVBoxLayout()
        page2.setLayout(layout)
        def chooseISO(event=None):
            source, _ = QFileDialog.getOpenFileName(
                caption = self.tr("Choose an image to clone"),
                filter = self.tr("ISO Images (*.iso);;All files (*)")
            )
            page2.setField("source", source)
        chooseButton=QPushButton("Choose the source of the core system")
        chooseButton.clicked.connect(chooseISO)
        layout.addWidget(chooseButton)
        layout.addWidget(sourceEdit)
        def page2validate(foo=None):
            command="/usr/bin/python3"
            source=page2.field("source")
            iso_filesize = None
            if source.startswith("AUTOCLONE"):
                source="own"
                # now calculate iso_filessize from mountpoint data
                mountpoint = glob.glob("/usr/lib/live/mount/persistence/sd*")[0]
                sourceDevice=mountpoint.replace("/usr/lib/live/mount/persistence","/dev")[:-1]
                cmd = f"""sudo LANG=C /usr/sbin/sfdisk -l {sourceDevice} | grep '{sourceDevice}1'""".format(sourceDevice=sourceDevice)
                reply, _ = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
                m = re.match(r".* ([.\d]+)G.*", reply.decode())
                size=float(m.group(1)) # size in GB
                iso_filesize = size*1024*1024*1024

            for device in self.field("selectedDisks").split(","):
                rindex=mw.rindexFromDev(device)
                if rindex==None:
                    return True
                if iso_filesize != None:
                    args=[
                        "{}/makeLiveStick.py".format(mw.wd),
                        "--vfat-size",
                        vfatLineEdit.text(),
                        "--no-persistence-seed",
                        "--iso-filesize",
                        str(int(iso_filesize/(1024**3))),
                        "--source",
                        source,
                        device
                    ]
                else:
                    args=[
                        "{}/makeLiveStick.py".format(mw.wd),
                        "--vfat-size",
                        vfatLineEdit.text(),
                        "--no-persistence-seed",
                        "--source",
                        source,
                        device
                    ]
                ### add a new tab and a monitor
                n=mw.ui.tabWidget.count()
                title=self.tr("Cloning to {}").format(device)
                mw.ui.tabWidget.addTab(
                    QWidget(),
                    self.tr("Cloning to {}").format(device)
                )
                m=CloneMonitor(command, args, mw, n, rindex, iso_filesize = iso_filesize, environ=os.environ)
                mw.monitors[m]=device
                m.start()
            return True
        page2.validatePage=page2validate
        self.addPage(page2)
        return

class DiskinstallWizard(QWizard):
    """
    The wizard associated to the "Clone" button of live-clone
    but only when the --diskinstall option is triggered
    """

    def __init__(self, mw):
        """
        the constructor
        @param mw a MyMain (<= QMainWindow) instance
        """
        QWizard.__init__(self, mw)
        self.setWindowTitle(self.tr("Clone a Freeduc system to the hard disk"))
        def page1validate(foo=None):
            source="own"
            device="/dev/sda"
            # now calculate iso_filessize from mountpoint data
            mountpoints = glob.glob("/usr/lib/live/mount/persistence/sd*")
            if not mountpoints:
                QMessageBox.warning(mw, self.tr("Warning"), self.tr("You are not running this application from a live USB key. Aborting the clone action."))
                return True
            mountpoint = mountpoints[0]
            sourceDevice=mountpoint.replace("/usr/lib/live/mount/persistence","/dev")[:-1]
            # now sourceDevice is something like /dev/sdb
            cmd = f"""sudo LANG=C /usr/sbin/sfdisk -l {sourceDevice} | grep '{sourceDevice}1'""".format(sourceDevice=sourceDevice)
            reply, _ = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
            m = re.match(r".* ([.\d]+)G.*", reply.decode())
            size=float(m.group(1)) # size in GB
            iso_filesize = size*1024*1024*1024
            ### add a new tab and a monitor
            title=self.tr("Cloning to {}").format(device)
            mw.tabWidget.addTab(
                QWidget(),
                self.tr("Cloning to {}").format(device)
            )
            command="/bin/sh"
            args = ["/usr/sbin/cloneToSda"]
            m=CloneMonitor(command, args, mw, 1, 0, iso_filesize = iso_filesize, environ=os.environ)
            mw.monitors[m]=device
            m.start()
            
            return True
        page1=QWizardPage(self)
        page1.setTitle(self.tr("DANGER ZONE"))
        page1.setSubTitle(self.tr("Your USB disk is /dev/{usb}, its image would be written on the hard disk /dev/{hard}. Are you really sure that you want to erase and loose any data available on the hard disk /dev/{hard}?").format(usb=mw.owndisk, hard="sda"))
        layout=QVBoxLayout()
        cb1 = QCheckBox(self.tr("Yes, I am sure that I want to erase the hard disk."))
        cb2 = QCheckBox(self.tr("Yes again, I am completely sure!"))
        page1.registerField("ok1", cb1)
        page1.registerField("ok2", cb2)
        layout.addWidget(cb1)
        layout.addWidget(cb2)
        page1.setLayout(layout)
        page1.validatePage=page1validate
        self.addPage(page1)
        return

class RunWizard(QWizard):
        """
        wizard for the runButton of live-clone
        """
        def __init__(self, mw, devices, available):
            """
            the constructor
            @param mw a MyMain (<= QMainWindow) instance
            @param devices a list of paths to USB devices
            @param available an Available instance
            """
            QWizard.__init__(self, mw)
            disks=sorted([d for d in devices])
            self.setWindowTitle(self.tr("Run a disk in a virtual machine"))
            page1=QWizardPage(self)
            page1.setTitle(self.tr("Select one disk"))
            page1.setSubTitle(self.tr("Which disk do you want to launch in a virtual machine?"))
            layout=QVBoxLayout()
            radios={}
            first=True
            for d in disks:
                rb=QRadioButton(d, parent=page1)
                if first:
                    rb.setChecked(True)
                    first=False
                layout.addWidget(rb)
                radios[rb]=d
            disks=QLineEdit(page1)
            disks.setReadOnly(True)
            disks.hide()
            def page1validate(foo=None):
                foundDisks=[]
                for rb in radios:
                    if rb.isChecked():
                        foundDisks.append(radios[rb])
                disks.setText(",".join(foundDisks))
                return bool(foundDisks)
            page1.validatePage=page1validate
            page1.registerField("disks", disks)
            page1.setLayout(layout)
            self.addPage(page1)
            return