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

from anacondalib import INSTALLER_VM_MEMORY_MB, VirtInstallMachineCase, disk_images, run_boot
from installer import Installer
from review import Review
from storage import Storage
from storagelib import StorageCase  # pylint: disable=import-error
from testlib import test_main, timeout  # pylint: disable=import-error


class TestStorageCockpitIntegration_E2E(VirtInstallMachineCase, StorageCase):
    provision = {"machine1": {"memory_mb": INSTALLER_VM_MEMORY_MB}}

    def cockpitCreateBootloaderPartitions(self):
        self.click_dropdown(self.card_row("Storage", 1), "Create partition table")
        self.confirm()

        self.click_dropdown(self.card_row("Storage", 2), "Create partition")
        if self.is_efi:
            self.dialog({"size": 100, "type": "efi", "mount_point": "/boot/efi"})
        else:
            self.dialog({"size": 1, "type": "biosboot"})

        self.click_dropdown(self.card_row("Storage", 3), "Create partition")
        self.dialog({"size": 1070, "type": "ext4", "mount_point": "/boot"})

    @disk_images([('', 15), ('', 15), ('', 15)])
    @timeout(1200)
    def testRAIDonPartitions(self):
        """
        Description:
            Test scenario with three disks:

            - ESP and '/boot' are on 'vda'.
            - A RAID 0 array is created using partitions on 'vdb1' and 'vdc1'.
            - The root filesystem ('/') is placed on the RAID 0 device.

        Expected results:
            - The installer can create a RAID0 array on partitions and use it for rootfs
            - The installed system storage layout is as expected
        """

        b = self.browser
        m = self.machine
        i = Installer(b, m, scenario="use-configured-storage")
        s = Storage(b, m)
        r = Review(b, m)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)
        s.select_disks([("vdb", True), ("vda", True), ("vdc", True)])

        s.enter_cockpit_storage()

        # Create ESP and /boot partitions on vda
        self.cockpitCreateBootloaderPartitions()

        self.click_dropdown(self.card_row("Storage", 4), "Create partition")
        self.dialog({ "type": "empty"})

        # Create partitions for the RAID device on vdb and vdc
        self.click_dropdown(self.card_row("Storage", 5), "Create partition table")
        self.dialog({"type": "gpt"})
        self.click_dropdown(self.card_row("Storage", 6), "Create partition")
        self.dialog({"type": "empty"})

        self.click_dropdown(self.card_row("Storage", 7), "Create partition table")
        self.dialog({"type": "gpt"})
        self.click_dropdown(self.card_row("Storage", 8), "Create partition")
        self.dialog({"type": "empty"})

        # Create RAID1 on vda
        self.click_dropdown(self.card_header("Storage"), "Create MDRAID device")
        self.dialog_wait_open()
        self.dialog_set_val("level", "raid0")
        self.dialog_set_val("disks", {"/dev/vda3": True, "/dev/vdb1": True, "/dev/vdc1": True})
        self.dialog_set_val("name", "SOMERAID")
        self.dialog_apply()
        self.dialog_wait_close()
        b.wait_visible(self.card_row("Storage", name="/dev/md/SOMERAID"))

        # Format and mount the RAID device partition as /
        self.click_dropdown(self.card_row("Storage", 9), "Format")
        self.dialog({"type": "ext4", "mount_point": "/"})

        # Exit the cockpit-storage iframe and return to installation
        b.switch_to_top()
        s.return_to_installation()
        s.return_to_installation_confirm()

        i.reach(i.steps.REVIEW)

        def checkStorageReview(prefix=""):
            with b.wait_timeout(30):
                for disk in ("vda", "vdb", "vdc"):
                    r.check_disk(disk, f"16.1 GB {disk} (Virtio Block Device)")
                    r.check_disk_row(disk, "/", "vda3, vdc1, vdb1, RAID", "47.1 GB", False, prefix=prefix)

        # verify review screen
        checkStorageReview()
        r.check_checkbox_not_present()
        self.install(needs_confirmation=False)

        # 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")
        vda3 = next(part for part in vda["children"] if part["name"] == "vda3")
        self.assertIn("raid0", [part["type"] for part in vda3["children"]])

        vdb = next(dev for dev in block_devs if dev["name"] == "vdb")
        vdb1 = next(part for part in vdb["children"] if part["name"] == "vdb1")
        self.assertIn("raid0", [part["type"] for part in vdb1["children"]])

        vdc = next(dev for dev in block_devs if dev["name"] == "vdc")
        vdc1 = next(part for part in vdc["children"] if part["name"] == "vdc1")
        self.assertIn("raid0", [part["type"] for part in vdc1["children"]])

        raid_root = next(part for part in vda3["children"] if part["type"] == "raid0")
        self.assertEqual(raid_root["mountpoints"], ["/"])
        self.assertEqual(raid_root["size"], "43.9G")

    @disk_images([("", 15), ("", 15), ("", 15)])
    @run_boot("bios", "efi")
    @timeout(1200)
    def testRAID0Scenario_0(self):
        """
        Description:
            Test scenario with three disks:

            - ESP or BIOS boot and '/boot' are on 'vda'.
            - A RAID 0 array is created using 'vdb' and 'vdc'.
            - The root filesystem ('/') is placed on the RAID 0 device.

        Expected results:
            - The installer can create a RAID 0 array on disks and use it for rootfs.
            - The installed system storage layout is as expected.
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m, scenario="use-configured-storage")
        s = Storage(b, m)
        r = Review(b, m)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)
        s.select_disks([("vdb", True), ("vda", True), ("vdc", True)])

        s.enter_cockpit_storage()

        # Create RAID device on vdb, and vdc
        self.click_dropdown(self.card_header("Storage"), "Create MDRAID device")
        self.dialog_wait_open()
        self.dialog_set_val("level", "raid0")
        self.dialog_set_val("disks", {"/dev/vdb": True, "/dev/vdc": True})
        self.dialog_set_val("name", "SOMERAID")
        self.dialog_apply()
        self.dialog_wait_close()

        # Create bootloader and /boot partitions on vda
        if self.is_efi:
            self.cockpitCreateBootloaderPartitions()
        else:
            self.click_dropdown(self.card_row("Storage", 1), "Create partition table")
            self.confirm()
            self.click_dropdown(self.card_row("Storage", 2), "Create partition")
            self.dialog({"size": 1, "type": "biosboot"})
            self.click_dropdown(self.card_row("Storage", 3), "Create partition")
            self.dialog({"size": 1070, "type": "ext4", "mount_point": "/boot"})

        # Create / partition on the RAID device
        self.click_dropdown(self.card_row("Storage", 7), "Create partition table")
        self.dialog({"type": "gpt"})
        self.click_dropdown(self.card_row("Storage", 8), "Create partition")
        self.dialog({"type": "ext4", "mount_point": "/"})

        def checkStorageReview(prefix=""):
            with b.wait_timeout(30):
                disk = "vda"
                r.check_disk(disk, f"16.1 GB {disk} (Virtio Block Device)", prefix=prefix)
                if self.is_efi:
                    r.check_disk_row(disk, "/boot/efi", "vda1", "99.6 MB", False, prefix=prefix)
                r.check_disk_row(disk, "/boot", "vda2", "1.07 GB", False, prefix=prefix)

                disk = "MDRAID-SOMERAID"
                r.check_disk(disk, f"32.2 GB {disk} (MDRAID set (stripe))", prefix=prefix)
                r.check_disk_row(disk, "/", "SOMERAID1, RAID", "32.2 GB", False, prefix=prefix)

        # Exit the cockpit-storage iframe and return to installation
        b.switch_to_top()
        s.return_to_installation()
        checkStorageReview(prefix="#cockpit-storage-integration-check-storage-dialog")
        s.return_to_installation_confirm()

        i.reach(i.steps.REVIEW)

        # verify review screen
        checkStorageReview()
        r.check_checkbox_not_present()
        self.install(needs_confirmation=False)

        # 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")
        vda2 = next(part for part in vda["children"] if part["name"] == "vda2")
        self.assertEqual(vda2["mountpoints"], ["/boot"])

        vdb = next(dev for dev in block_devs if dev["name"] == "vdb")
        raid_dev = next(part for part in vdb["children"] if part["type"] == "raid0")
        raid_part = next(part for part in raid_dev["children"] if part["type"] == "part")
        self.assertEqual(raid_part["mountpoints"], ["/"])
        self.assertEqual(raid_part["size"], "30G")

    @timeout(1200)
    def testSWAP(self):
        """
        Description:
            Test an installation using 'storage editor' and including a disk-based swap partition

        Expected results:
            - The install process should complete successfully and the installed system should boot as expected.
            - The expected partition layout should be created on the target device(s).
            - The disk-based swap partition should be active (you can check /proc/swaps or the output of swapon --show).
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m, scenario="use-configured-storage")
        s = Storage(b, m)
        r = Review(b, m)

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

        s.enter_cockpit_storage()
        # Create ESP and /boot partitions on vda
        self.cockpitCreateBootloaderPartitions()
        # Create a partition for /
        self.click_dropdown(self.card_row("Storage", 4), "Create partition")
        self.dialog({"size": 4096, "type": "ext4", "mount_point": "/"})
        # Create a partition for swap
        self.click_dropdown(self.card_row("Storage", 5), "Create partition")
        self.dialog({"size": 1024, "type": "swap"})
        # Exit the cockpit-storage iframe and return to installation
        s.exit_cockpit_storage()
        s.check_scenario_selected("use-configured-storage")

        i.reach(i.steps.REVIEW)
        r.check_disk("vda", "16.1 GB vda (Virtio Block Device)")
        r.check_disk_row("vda", "/boot", "vda2", "1.07 GB", False)
        r.check_disk_row("vda", "/", "vda3", "4.10 GB", False)
        r.check_disk_row("vda", "swap", "vda4", "1.02 GB", False)
        self.install(needs_confirmation=False)

        # Check that the expected partition layout is created
        lsblk = s.get_lsblk_json()
        block_devs = lsblk["blockdevices"]
        vda = next(dev for dev in block_devs if dev["name"] == "vda")
        for dev, mountpoint in (
            ("vda2", "/boot"),
            ("vda3", "/"),
            ("vda4", "[SWAP]"),
        ):
            part = next(part for part in vda["children"] if part["name"] == dev)
            self.assertEqual(part["mountpoints"], [mountpoint])

        # Check that the swap partition is active
        self.assertIn("/dev/vda4", m.execute("swapon --show"))


if __name__ == '__main__':
    test_main()
