#!/usr/bin/python3
#
# Copyright (C) 2024 Red Hat, Inc.
# SPDX-License-Identifier: LGPL-2.1-or-later

import os

from anacondalib import INSTALLER_VM_MEMORY_MB, VirtInstallMachineCase, disk_images
from installer import Installer
from password import Password
from progress import Progress
from review import Review
from storage import Storage
from testlib import test_main, timeout  # pylint: disable=import-error
from utils import get_pretty_name

TEST_DIR = os.path.dirname(__file__)
ROOT_DIR = os.path.dirname(TEST_DIR)
BOTS_DIR = f'{ROOT_DIR}/bots'


class TestStorageEraseAll_E2E(VirtInstallMachineCase):
    provision = {"machine1": {"memory_mb": INSTALLER_VM_MEMORY_MB}}

    @timeout(1200)
    def testGuidedEncrypted(self):
        """
        Description:
            Test installation with erase-all and encryption enabled

        Expected results:
            - The installed system is bootable and the encryption is enabled
        """
        b = self.browser
        i = Installer(b, self.machine)
        r = Review(b, self.machine)
        s = Storage(b, self.machine)
        p = Password(b, s.encryption_id_prefix)
        pro = Progress(b)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)

        # Check the default mode
        s.check_scenario_selected("erase-all")

        i.next()
        # Disk Encryption
        s.set_encryption_selected(True)
        strong_password = "Rwce82ybF7dXtCzFumanchu!!!!!!!!"
        p.set_password(strong_password)
        p.set_password_confirm(strong_password)

        # Reboot and check installed system
        i.reach(i.steps.REVIEW)
        r.check_checkbox_not_present()
        i.begin_installation(needs_confirmation=False)
        pro.wait_done()
        s.unlock_storage_on_boot(strong_password)
        self.handleReboot()
        s.check_post_install_encryption_enabled()

        # Check that the expected partition layout is created on the selected device
        lsblk = s.get_lsblk_json()
        block_devs = lsblk["blockdevices"]
        vda = next(dev for dev in block_devs if dev["name"] == "vda")
        vda3 = next(part for part in vda["children"] if part["name"] == "vda3")
        vda3_luks = next(part for part in vda3["children"] if part["type"] == "crypt")
        vda3_root = next(part for part in vda3_luks["children"] if "/" in part["mountpoints"])
        self.assertEqual(vda3_root["size"], "12.4G")

    @disk_images([("", 15), ("debian-testing", 15)])
    @timeout(1200)
    def testGuidedMultiSelect(self):
        """
        Description:
            Test dual boot installation with 'share with other OS' and multiple disks selected

        Expected results:
            - Both the new and existing OS are bootable and the partition layout is correct
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m)
        s = Storage(b, m)
        r = Review(b, m)
        p = Progress(b)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)

        s.check_disk_selected("vda", True)
        s.check_disk_selected("vdb", False)

        i.reach(i.steps.REVIEW)
        r.check_checkbox_not_present()
        i.begin_installation(needs_confirmation=False)
        p.wait_done()

        self.handleReboot()

        # Expect the new OS is the default grub entry
        pretty_name = get_pretty_name(m)
        self.assertIn("Fedora Linux", pretty_name)

        # Check that the expected partition layout is created on the selected device
        lsblk = s.get_lsblk_json()
        block_devs = lsblk["blockdevices"]
        vda = next(dev for dev in block_devs if dev["name"] == "vda")
        vda3 = next(part for part in vda["children"] if part["name"] == "vda3")
        vda3_root = next(part for part in vda3["children"] if "/" in part["mountpoints"])
        self.assertEqual(vda3_root["size"], "12.4G")

        vdb = next(dev for dev in block_devs if dev["name"] == "vdb")
        vdb_mountpoints = [part["mountpoints"] for part in vdb["children"]]
        self.assertEqual(vdb_mountpoints, [[]] * len(vdb_mountpoints))

        m.reboot()
        self.selectBootMenuEntry(2)
        # FIXME: Debian will not boot unless we remove the vda disk
        # pretty_name = get_pretty_name(m)
        # self.assertIn("Debian GNU/Linux", pretty_name)

    @disk_images([("", 15), ("", 15)])
    @timeout(1200)
    def testGuidedMultiEmptyAll(self):
        """
        Description:
            Test installation with erase-all and multiple empty disks selected

        Expected results:
            - The installed system is bootable and the partition layout is correct
        """

        b = self.browser
        m = self.machine
        i = Installer(b, self.machine)
        r = Review(b, m)
        s = Storage(b, self.machine)
        p = Progress(b)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)

        s.check_disk_selected("vda", True)
        s.select_disks([("vdb", True), ("vda", True)])

        i.reach(i.steps.REVIEW)
        r.check_checkbox_not_present()
        i.begin_installation(needs_confirmation=False)
        p.wait_done()

        self.handleReboot()

        # Expect the new OS is the default grub entry
        pretty_name = get_pretty_name(m)
        self.assertIn("Fedora Linux", pretty_name)

        # Check that the expected partition layout is created and spans all the target devices
        lsblk = s.get_lsblk_json()
        block_devs = lsblk["blockdevices"]
        vda = next(dev for dev in block_devs if dev["name"] == "vda")
        vdb = next(dev for dev in block_devs if dev["name"] == "vdb")
        vda3 = next(part for part in vda["children"] if part["name"] == "vda3")
        vdb1 = next(part for part in vdb["children"] if part["name"] == "vdb1")
        vda3_root = next(part for part in vda3["children"] if "/" in part["mountpoints"])
        vdb1_root = next(part for part in vdb1["children"] if "/" in part["mountpoints"])
        self.assertEqual(vda3_root["size"], "15G")
        self.assertEqual(vdb1_root["size"], "15G")


if __name__ == '__main__':
    test_main()
