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

from anacondalib import VirtInstallMachineCase, disk_images
from installer import Installer
from language import Language
from review import Review
from storage import BOOT_PARTITION_DEFAULT_SIZE, EFI_PARTITION_DEFAULT_SIZE, EFI_PARTITION_DEFAULT_SIZE_MB, Storage
from testlib import (
    nondestructive,
    test_main,
)


@nondestructive
class TestStorageReclaimSpace(VirtInstallMachineCase):
    def setup_partitions(self, s):
        disk = "/dev/vda"
        btrfsname = "btrfstest"
        s.partition_disk(
            disk,
            [
                (EFI_PARTITION_DEFAULT_SIZE, "efi"),
                ("6GiB", "ext4"),
                ("5GiB", "btrfs", "-f", "-L", f"{btrfsname}A"),
                ("", "btrfs", "-f", "-L", f"{btrfsname}B"),
            ],
        )

    def _testReclaimSpaceOptional_partition_disk(self):
        b = self.browser
        m = self.machine
        s = Storage(b, m)

        disk = "/dev/vda"
        btrfsname = "btrfstest"
        s.partition_disk(
            disk,
            [
                (EFI_PARTITION_DEFAULT_SIZE, "efi"),
                (BOOT_PARTITION_DEFAULT_SIZE, "ext4"),
                ("5GiB", "btrfs", "-f", "-L", f"{btrfsname}A"),
            ],
        )

    def testReclaimSpaceOptional(self):
        """
        Description:
            Test the reclaim space scenario when there is enough free space
            for the installation.

        Expected results:
            - The reclaim space checkbox is not selected by default.
            - The reclaim space checkbox is enabled.
        """

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

        dev = "vda"

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

        s.set_scenario("use-free-space")
        s.reclaim_check_checkbox(False, False)

        s.reclaim_set_checkbox(True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        # Test cancelled reclaim is not taking effect
        s.reclaim_remove_device("vda2")
        s.reclaim_check_action_present("vda2", "delete")
        s.reclaim_modal_cancel()

        # Continue without reclaiming space
        s.reclaim_check_checkbox(True, False)
        s.reclaim_set_checkbox(False)

        i.reach(i.steps.REVIEW)

        r.check_disk_row(dev, "/", "vda5, LVM", "5.91 GB", True, "xfs")

    def _testReclaimSpaceDeleteBtrfsSubvolumes_partition_disk(self):
        b = self.browser
        m = self.machine
        s = Storage(b, m)

        self.setup_partitions(s)

    def testReclaimSpaceDeleteBtrfsSubvolumes(self):
        """
        Description:
            Test the reclaim space scenario with deleting BTRFS subvolumes

        Expected results:
            - The BTRFS subvolumes can be deleted in the reclaim space dialog.
            - Deleting a parent partition will delete all contained partitions.
            - The delete operation can be undone.
            - The installation can proceed after deleting
            - The installation can proceed after allocating enough space by deleting partitions.
        """

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

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

        s.set_scenario("use-free-space")
        s.reclaim_check_checkbox(True, True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        s.reclaim_check_available_space("2.06 MB")

        # Check that all partitions are present
        s.reclaim_check_device_row("vda", deviceType="disk", space="16.1 GB")
        s.reclaim_check_device_row("vda1", deviceType="efi", space=EFI_PARTITION_DEFAULT_SIZE_MB)
        s.reclaim_check_device_row("vda2", deviceType="ext4", space="6.44 GB")
        s.reclaim_check_device_row("vda3", deviceType="btrfs", space="5.37 GB")
        s.reclaim_check_device_row("vda4", deviceType="btrfs", space="3.77 GB")

        # Check that deleting a disk will delete all contained partitions
        s.reclaim_remove_device("vda")
        for device in ["vda1", "vda2", "vda3"]:
            s.reclaim_check_action_present(device, "delete")

        s.reclaim_check_available_space("16.1 GB")

        # Undo disk device deletion
        s.reclaim_undo_action("vda")
        for device in ["vda1", "vda2", "vda3"]:
            s.reclaim_check_action_present(device, "delete", present=False)

        s.reclaim_check_available_space("2.06 MB")

        # Check that actions for devices whose parents are marked for deletion are not sent to blivet
        s.reclaim_remove_device("vda4")
        s.reclaim_remove_device("vda")
        s.reclaim_modal_submit()

        i.back()
        s.reclaim_check_checkbox(True, True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        # Remove one partition and allocate enough space for the installation
        s.reclaim_remove_device("vda2")
        s.reclaim_check_action_present("vda2", "delete")
        s.reclaim_check_available_space("6.44 GB")

        # Check that deleting a parent disk which contains partitions marked for deletion
        # correctly calculates the available space
        s.reclaim_remove_device("vda")
        s.reclaim_check_available_space("16.1 GB")
        s.reclaim_undo_action("vda")
        s.reclaim_check_available_space("6.44 GB")

        s.reclaim_modal_submit()
        i.reach(i.steps.REVIEW)

        r.check_disk_row("vda", parent="vda2", action="delete")

    def _testReclaimSpaceShrinkBtrfsSubvolumes_partition_disk(self):
        b = self.browser
        m = self.machine
        s = Storage(b, m)

        self.setup_partitions(s)

    def testReclaimSpaceShrinkBtrfsSubvolumes(self):
        """
        Description:
            Test the reclaim space scenario with shrinking partitions

        Expected results:
            - The shrinking is not available for disks.
            - Shrinking BTRFS subvolumes is not available.
            - A warning is shown when trying to shrink a partition too much.
            - The shrinking operation can be undone.
            - The installation can proceed after allocating enough space by shrinking partitions.
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m)
        s = Storage(b, m)
        r = Review(b, m)

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

        s.set_scenario("use-free-space")
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        s.reclaim_check_available_space("2.06 MB")

        # Check that shrinking action is only available for partitions
        s.reclaim_check_action_button_present("vda", "shrink", present=False)
        for device in ["vda2"]:
            s.reclaim_check_action_button_present(device, "shrink")

        for device in ["vda3", "vda4"]:
            # Shrinking Btrfs partitions is not yet supported by blivet:
            # https://bugzilla.redhat.com/show_bug.cgi?id=962143
            s.reclaim_check_action_button_present(device, "shrink", disabled=True)

        # Shrink a partition too much and expect a warning
        s.reclaim_shrink_device("vda2", "0.100", "6.44")
        s.reclaim_modal_submit_and_check_warning("Unable to schedule resizing of vda2")
        s.reclaim_undo_action("vda2")

        # Shrink a partition
        s.reclaim_shrink_device("vda2", "1", "6.44")
        s.reclaim_check_available_space("5.44 GB")

        # Undo the shrink actions
        for device in ["vda2"]:
            s.reclaim_undo_action(device)
            s.reclaim_check_action_present(device, "shrink", present=False)

        s.reclaim_check_available_space("2.06 MB")

        # Shrink one partition and allocate enough space for the installation
        s.reclaim_shrink_device("vda2", "1", "6.44")

        s.reclaim_modal_submit()
        i.reach(i.steps.REVIEW)

        r.check_disk_row("vda", parent="vda2", size="1.00 GB", action="resized from 6.44 GB")

    def _testReclaimExt4onLUKS_partition_disk(self):
        b = self.browser
        m = self.machine
        s = Storage(b, m)

        disk = "/dev/vda"
        s.partition_disk(disk, [("4GiB", "ext4")])
        s.create_luks_partition(f"{disk}1", "einszweidrei", "encrypted-vol")

    def testReclaimExt4onLUKS(self):
        """
        Description:
            Test the reclaim space scenario with LUKS partitions

        Expected results:
            - The LUKS partition cannot be shrunk
            - The LUKS partition can be deleted
            - The installation can proceed after deleting the LUKS partition - password is not needed
        """
        # Shrinking LUKS partitions is not yet supported
        # Deletion of locked partitions should be possible

        b = self.browser
        m = self.machine
        i = Installer(b, m)
        s = Storage(b, m)

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

        s.set_scenario("use-free-space")
        s.reclaim_set_checkbox(True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        b.wait_visible("#reclaim-space-modal")
        b.assert_pixels(
            "#reclaim-space-modal",
            "reclaim-space-modal-encrypted",
            ignore=["#reclaim-space-modal-hint"],
        )

        s.reclaim_check_device_row("vda1", deviceType="luks", space="4.29 GB", locked=True)
        s.reclaim_check_action_button_present("vda1", "shrink", disabled=True)
        s.reclaim_check_action_button_present("vda1", "delete")

    @disk_images([("fedora-rawhide", 15)])
    def testDeletePartition(self):
        """
        Description:
            Test reclaiming space from an existing Fedora system by deleting partitions

        Expected results:
            - The review screen shows the affected system when partitions from an existing system are deleted
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m)
        s = Storage(b, m)
        r = Review(b, m)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)
        s.set_scenario("use-free-space")
        i.next(next_page=i.steps.INSTALLATION_METHOD)
        b.wait_visible("#reclaim-space-modal")

        b.assert_pixels(
            "#reclaim-space-modal",
            "reclaim-space-modal-fedora",
            ignore=["td[data-label=Space]", "#reclaim-space-modal-hint"],
        )

        s.reclaim_remove_device("vda3")
        s.reclaim_modal_submit()
        # Backend gives warning on the 100MB /boot/efi partition
        i.next(should_fail=True)

        i.reach(i.steps.REVIEW)
        r.check_affected_system("Fedora Linux", [("vda3", ["boot", "home", "root", "var"])])
        r.check_some_erased_checkbox_label()


    def _testReclaimSpaceWindows_partition_disk(self):
        b = self.browser
        m = self.machine
        s = Storage(b, m)

        disk = "/dev/vda"
        s.partition_disk(disk, [("1MiB", "efi"), ("1GiB", "microsoft-reserved"), ("10GiB", "basic-data"), ("500MiB", "microsoft-recovery")])

    def testReclaimSpaceWindows(self):
        """
        Description:
            Test reclaiming space from an existing Windows system by deleting partitions

        Expected results:
            - The Windows partitions are identified and displayed in the reclaim space dialog
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m)
        s = Storage(b, m)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)
        s.set_scenario("use-free-space")
        s.reclaim_set_checkbox(True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        b.wait_visible("#reclaim-space-modal")
        b.assert_pixels(
            "#reclaim-space-modal",
            "reclaim-space-modal-windows",
            ignore=["td[data-label=Space]", "#reclaim-space-modal-hint"],
        )

    @disk_images([("debian-testing", 15)])
    def testReclaimSpaceShrink(self):
        """
        Description:
            Test reclaiming space from an existing Debian system by resizing partitions

        Expected results:
            - The shrink operation works also with locales that use a comma as decimal separator
            - The review screen shows the affected system when partitions from an existing system are resized
        """
        b = self.browser
        m = self.machine
        i = Installer(b, m)
        l = Language(b, m)
        s = Storage(b, m)
        r = Review(b, m)

        i.open()
        i.reach(i.steps.INSTALLATION_METHOD)
        s.set_scenario("use-free-space")
        s.reclaim_set_checkbox(True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        s.reclaim_shrink_device("vda1", "7")
        s.reclaim_modal_submit()

        # existing efi partition is smaller than the recommended 500 MiB
        # so we need to click 'Next' twice to move past the storage configuration step
        i.next(next_page=i.steps.STORAGE_CONFIGURATION)

        i.reach(i.steps.REVIEW)
        # Don't specify the exact original size as this might change with image refreshes
        r.check_disk_row("vda", parent="vda1", size="7.00 GB", action="resized from 11.7 GB")
        r.check_resized_system("Debian", ["vda1"])

        i.reach_on_sidebar(i.steps.LANGUAGE)
        l.select_locale("cs_CZ", is_common=False)
        i.reach(i.steps.INSTALLATION_METHOD)
        s.reclaim_set_checkbox(True)
        i.next(next_page=i.steps.INSTALLATION_METHOD)

        s.reclaim_shrink_device("vda1", "7,10", ariaLabel="zmenšit")
        s.reclaim_modal_submit()
        # existing efi partition is smaller than the recommended 500 MiB
        # so we need to click 'Next' twice to move past the storage configuration step
        i.next(next_page=i.steps.STORAGE_CONFIGURATION)

        i.reach(i.steps.REVIEW)
        r.check_disk_row("vda", parent="vda1", size="7,10 GB", action="změněna velikost z 11,7 GB")


if __name__ == '__main__':
    test_main()
