/*
 * This file is part of the KubeVirt project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Copyright The KubeVirt Authors.
 *
 */

package storage

import (
	"context"
	"fmt"
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	k8sv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/rand"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"

	"kubevirt.io/kubevirt/pkg/apimachinery/patch"
	"kubevirt.io/kubevirt/pkg/libdv"
	"kubevirt.io/kubevirt/pkg/libvmi"
	backendstorage "kubevirt.io/kubevirt/pkg/storage/backend-storage"
	"kubevirt.io/kubevirt/pkg/storage/cbt"

	"kubevirt.io/kubevirt/tests/console"
	cd "kubevirt.io/kubevirt/tests/containerdisk"
	"kubevirt.io/kubevirt/tests/decorators"
	"kubevirt.io/kubevirt/tests/exec"
	"kubevirt.io/kubevirt/tests/framework/kubevirt"
	. "kubevirt.io/kubevirt/tests/framework/matcher"
	"kubevirt.io/kubevirt/tests/libmigration"
	"kubevirt.io/kubevirt/tests/libnamespace"
	"kubevirt.io/kubevirt/tests/libpod"
	"kubevirt.io/kubevirt/tests/libstorage"
	"kubevirt.io/kubevirt/tests/libvmops"
	"kubevirt.io/kubevirt/tests/libwait"
	"kubevirt.io/kubevirt/tests/testsuite"
)

var _ = Describe(SIG("CBT", func() {
	var (
		err        error
		virtClient kubecli.KubevirtClient
		vm         *v1.VirtualMachine
		vmi        *v1.VirtualMachineInstance
	)

	BeforeEach(func() {
		virtClient = kubevirt.Client()
	})

	It("VM matches cbt label selector, then unmatches", func() {
		vm = libstorage.RenderVMWithDataVolumeTemplate(libdv.NewDataVolume(
			libdv.WithRegistryURLSource(cd.DataVolumeImportUrlForContainerDisk(cd.ContainerDiskAlpine)),
			libdv.WithNamespace(testsuite.GetTestNamespace(nil)),
			libdv.WithStorage(),
		),
			libvmi.WithLabels(cbt.CBTLabel),
			libvmi.WithRunStrategy(v1.RunStrategyAlways),
		)
		volumeName := vm.Spec.Template.Spec.Volumes[0].Name

		By(fmt.Sprintf("Creating VM %s with CBT label", vm.Name))
		vm, err := virtClient.VirtualMachine(vm.Namespace).Create(context.Background(), vm, metav1.CreateOptions{})
		Expect(err).ShouldNot(HaveOccurred())
		libstorage.WaitForCBTEnabled(virtClient, vm.Namespace, vm.Name)
		vmi, err = virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
		Expect(err).ShouldNot(HaveOccurred())

		By("Verify CBT overlay exists")
		stdout := libpod.RunCommandOnVmiPod(vmi, []string{"find", cbt.PathForCBT(vmi), "-type", "f", "-name", fmt.Sprintf("%s.qcow2", volumeName)})
		Expect(stdout).To(ContainSubstring(cbt.GetQCOW2OverlayPath(vmi, volumeName)))

		By("Remove CBT Label")
		vm, err = virtClient.VirtualMachine(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
		Expect(err).ShouldNot(HaveOccurred())
		delete(vm.Labels, cbt.CBTKey)
		patch, err := patch.New(patch.WithAdd("/metadata/labels", vm.Labels)).GeneratePayload()
		Expect(err).ToNot(HaveOccurred())

		vm, err = virtClient.VirtualMachine(vm.Namespace).Patch(context.Background(), vm.Name, types.JSONPatchType, patch, metav1.PatchOptions{})
		Expect(err).ToNot(HaveOccurred())

		By("Verify CBT state PendingRestart")
		Eventually(func() v1.ChangedBlockTrackingState {
			vm, err = virtClient.VirtualMachine(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
			Expect(err).ShouldNot(HaveOccurred())
			return cbt.CBTState(vm.Status.ChangedBlockTracking)
		}, 1*time.Minute, 3*time.Second).Should(Equal(v1.ChangedBlockTrackingPendingRestart))

		By("Restarting the VM")
		err = virtClient.VirtualMachine(vm.Namespace).Restart(context.Background(), vm.Name, &v1.RestartOptions{})
		Expect(err).ToNot(HaveOccurred())
		Eventually(ThisVM(vm)).WithTimeout(300 * time.Second).WithPolling(time.Second).Should(BeReady())

		By("Verify CBT state disabled")
		Eventually(func() v1.ChangedBlockTrackingState {
			vm, err = virtClient.VirtualMachine(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
			if err != nil {
				return v1.ChangedBlockTrackingUndefined
			}
			return cbt.CBTState(vm.Status.ChangedBlockTracking)
		}, 1*time.Minute, 3*time.Second).Should(Equal(v1.ChangedBlockTrackingDisabled))

		Eventually(func() v1.ChangedBlockTrackingState {
			vmi, err = virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
			if err != nil {
				return v1.ChangedBlockTrackingUndefined
			}
			return cbt.CBTState(vmi.Status.ChangedBlockTracking)
		}, 1*time.Minute, 3*time.Second).Should(Equal(v1.ChangedBlockTrackingDisabled))

		By("Verify CBT overlay deleted")
		vmi, err = virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
		Expect(err).ShouldNot(HaveOccurred())
		libwait.WaitUntilVMIReady(vmi, console.LoginToAlpine)
		pod, err := libpod.GetPodByVirtualMachineInstance(vmi, vmi.Namespace)
		Expect(err).ToNot(HaveOccurred())
		Expect(pod).NotTo(BeNil())

		var cmdOutput string
		cmdOutput, err = exec.ExecuteCommandOnPod(pod, "compute", []string{"ls", "-d", cbt.PathForCBT(vmi)})
		Expect(err.Error()).To(ContainSubstring("No such file or directory"))
		Expect(cmdOutput).To(BeEmpty())
	})

	DescribeTable("Patch to match cbt label selector", func(patchFunc func(vm *v1.VirtualMachine)) {
		vm = libstorage.RenderVMWithDataVolumeTemplate(libdv.NewDataVolume(
			libdv.WithRegistryURLSource(cd.DataVolumeImportUrlForContainerDisk(cd.ContainerDiskAlpine)),
			libdv.WithNamespace(testsuite.GetTestNamespace(nil)),
			libdv.WithStorage(),
		),
			libvmi.WithRunStrategy(v1.RunStrategyAlways),
		)
		volumeName := vm.Spec.Template.Spec.Volumes[0].Name

		By(fmt.Sprintf("Creating VM %s", vm.Name))
		_, err = virtClient.VirtualMachine(vm.Namespace).Create(context.Background(), vm, metav1.CreateOptions{})
		Expect(err).ToNot(HaveOccurred())
		Eventually(ThisVMIWith(vm.Namespace, vm.Name), 360).Should(BeInPhase(v1.Running))
		vm, err = virtClient.VirtualMachine(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
		Expect(err).ShouldNot(HaveOccurred())
		Expect(cbt.CBTState(vm.Status.ChangedBlockTracking)).To(Equal(v1.ChangedBlockTrackingUndefined))

		patchFunc(vm)

		Eventually(func() v1.ChangedBlockTrackingState {
			vm, err = virtClient.VirtualMachine(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
			Expect(err).ShouldNot(HaveOccurred())
			return cbt.CBTState(vm.Status.ChangedBlockTracking)
		}, 1*time.Minute, 3*time.Second).Should(Equal(v1.ChangedBlockTrackingPendingRestart))

		By("Restarting the VM")
		err = virtClient.VirtualMachine(vm.Namespace).Restart(context.Background(), vm.Name, &v1.RestartOptions{})
		Expect(err).ToNot(HaveOccurred())

		libstorage.WaitForCBTEnabled(virtClient, vm.Namespace, vm.Name)
		vmi, err = virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
		Expect(err).ShouldNot(HaveOccurred())

		stdout := libpod.RunCommandOnVmiPod(vmi, []string{"find", cbt.PathForCBT(vmi), "-type", "f", "-name", fmt.Sprintf("%s.qcow2", volumeName)})
		Expect(stdout).To(ContainSubstring(cbt.GetQCOW2OverlayPath(vmi, volumeName)))
	},
		Entry("patch vm", func(vm *v1.VirtualMachine) {
			patch, err := patch.New(patch.WithAdd("/metadata/labels", cbt.CBTLabel)).GeneratePayload()
			Expect(err).ToNot(HaveOccurred())

			_, err = virtClient.VirtualMachine(vm.Namespace).Patch(context.Background(), vm.Name, types.JSONPatchType, patch, metav1.PatchOptions{})
			Expect(err).ToNot(HaveOccurred())

		}),
		Entry("patch vm namespace", func(vm *v1.VirtualMachine) {
			Expect(libnamespace.AddLabelToNamespace(virtClient, vm.Namespace, cbt.CBTKey, "true")).ToNot(HaveOccurred())
		}),
	)

	DescribeTable("should create CBT overlay for hotplug volume and remove it on unplug",
		func(volumeMode k8sv1.PersistentVolumeMode) {
			var sc string
			var found bool
			if volumeMode == k8sv1.PersistentVolumeBlock {
				sc, found = libstorage.GetRWOBlockStorageClass()
			} else {
				sc, found = libstorage.GetRWOFileSystemStorageClass()
			}
			Expect(found).To(BeTrue(), "Expected storage class to be available for volume mode %s", volumeMode)

			By("Creating VM with CBT label")
			vm = libstorage.RenderVMWithDataVolumeTemplate(libdv.NewDataVolume(
				libdv.WithRegistryURLSource(cd.DataVolumeImportUrlForContainerDisk(cd.ContainerDiskAlpine)),
				libdv.WithNamespace(testsuite.GetTestNamespace(nil)),
				libdv.WithStorage(),
			),
				libvmi.WithLabels(cbt.CBTLabel),
				libvmi.WithRunStrategy(v1.RunStrategyAlways),
			)

			vm, err = virtClient.VirtualMachine(vm.Namespace).Create(context.Background(), vm, metav1.CreateOptions{})
			Expect(err).ShouldNot(HaveOccurred())
			Eventually(ThisVM(vm)).WithTimeout(300 * time.Second).WithPolling(time.Second).Should(BeReady())

			By("Waiting for CBT to be enabled")
			libstorage.WaitForCBTEnabled(virtClient, vm.Namespace, vm.Name)
			vmi, err = virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
			Expect(err).ShouldNot(HaveOccurred())
			vmi = libwait.WaitUntilVMIReady(vmi, console.LoginToAlpine)

			By(fmt.Sprintf("Creating DataVolume for hotplug with volume mode %s", volumeMode))
			hotplugVolumeName := "hotplug-vol"
			hotplugDV := libdv.NewDataVolume(
				libdv.WithBlankImageSource(),
				libdv.WithStorage(
					libdv.StorageWithStorageClass(sc),
					libdv.StorageWithVolumeMode(volumeMode),
					libdv.StorageWithVolumeSize(cd.BlankVolumeSize),
				),
			)
			hotplugDV, err = virtClient.CdiClient().CdiV1beta1().DataVolumes(vm.Namespace).Create(context.Background(), hotplugDV, metav1.CreateOptions{})
			Expect(err).ToNot(HaveOccurred())

			By("Hotplugging the DataVolume")
			vm = libstorage.AddHotplugDiskAndVolume(virtClient, vm, hotplugVolumeName, hotplugDV.Name)

			By("Waiting for hotplug volume to be ready")
			libstorage.WaitForHotplugToComplete(virtClient, vm, hotplugVolumeName, hotplugDV.Name, true)

			By("Verifying QCOW2 overlay was created for hotplug volume")
			vmi, err = virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
			Expect(err).ShouldNot(HaveOccurred())
			overlayPath := cbt.GetQCOW2OverlayPath(vmi, hotplugVolumeName)
			stdout := libpod.RunCommandOnVmiPod(vmi, []string{"find", cbt.PathForCBT(vmi), "-type", "f", "-name", fmt.Sprintf("%s.qcow2", hotplugVolumeName)})
			Expect(stdout).To(ContainSubstring(overlayPath), "Expected CBT overlay to exist for hotplug volume")

			By("Removing the hotplug volume")
			vm = libstorage.RemoveHotplugDiskAndVolume(virtClient, vm, hotplugVolumeName)

			By("Waiting for hotplug volume to be removed from VMI")
			libstorage.WaitForHotplugToComplete(virtClient, vm, hotplugVolumeName, hotplugDV.Name, false)

			By("Verifying QCOW2 overlay was deleted for unplugged volume")
			Eventually(func() string {
				return libpod.RunCommandOnVmiPod(vmi, []string{"find", cbt.PathForCBT(vmi), "-type", "f", "-name", fmt.Sprintf("%s.qcow2", hotplugVolumeName)})
			}, 60*time.Second, 2*time.Second).Should(BeEmpty(), "CBT overlay should be deleted for unplugged volume")
		},
		Entry("Filesystem", k8sv1.PersistentVolumeFilesystem),
		Entry("Block", decorators.RequiresBlockStorage, k8sv1.PersistentVolumeBlock),
	)

	It("should persist CBT data across restart", func() {
		runCBTPersistenceTest(virtClient, "restart")
	})

}))

var _ = Describe("[sig-compute-migrations] CBT with migration", decorators.SigComputeMigrations, func() {
	var virtClient kubecli.KubevirtClient

	BeforeEach(func() {
		virtClient = kubevirt.Client()
	})

	// NOTE: Currently there is a bug in libvirt where if the qcow2 overlay is on a RWO block storage,
	// the bitmap migration fails. resulting in the deletion of previous bitmaps. Will add a test with RWO block storage once the bug is fixed.
	// Bug: https://issues.redhat.com/browse/RHEL-145770
	It("should persist CBT data across live migration", decorators.RequiresTwoSchedulableNodes, decorators.RequiresRWXFsVMStateStorageClass, decorators.NoFlakeCheck, func() {
		runCBTPersistenceTest(virtClient, "migrate")
	})
})

func runCBTPersistenceTest(virtClient kubecli.KubevirtClient, op string) {
	sc, foundSC := libstorage.GetRWOFileSystemStorageClass()
	accessMode := k8sv1.ReadWriteOnce
	if op == "migrate" {
		sc, foundSC = libstorage.GetRWXFileSystemStorageClass()
		accessMode = k8sv1.ReadWriteMany
	}
	if !foundSC {
		Fail(fmt.Sprintf("Fail test when %s storage is not present", accessMode))
	}
	vm := libstorage.RenderVMWithDataVolumeTemplate(libdv.NewDataVolume(
		libdv.WithRegistryURLSource(cd.DataVolumeImportUrlForContainerDisk(cd.ContainerDiskAlpine)),
		libdv.WithStorage(
			libdv.StorageWithStorageClass(sc),
			libdv.StorageWithVolumeMode(k8sv1.PersistentVolumeFilesystem),
			libdv.StorageWithAccessMode(accessMode),
			libdv.StorageWithVolumeSize(cd.ContainerDiskSizeBySourceURL(cd.DataVolumeImportUrlForContainerDisk(cd.ContainerDiskAlpine))),
		),
	),
		libvmi.WithLabels(cbt.CBTLabel),
		libvmi.WithRunStrategy(v1.RunStrategyAlways),
	)
	volumeName := vm.Spec.Template.Spec.Volumes[0].Name

	vm, err := virtClient.VirtualMachine(testsuite.GetTestNamespace(vm)).Create(context.Background(), vm, metav1.CreateOptions{})
	Expect(err).ToNot(HaveOccurred())
	Eventually(ThisVM(vm)).WithTimeout(300 * time.Second).WithPolling(time.Second).Should(BeReady())
	vmi, err := virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
	Expect(err).ToNot(HaveOccurred())
	vmi = libwait.WaitUntilVMIReady(vmi, console.LoginToAlpine)

	By("Verifying CBT is enabled")
	libstorage.WaitForCBTEnabled(virtClient, vm.Namespace, vm.Name)

	By("Writing data to disk to trigger CBT tracking")
	Expect(console.RunCommand(vmi, "dd if=/dev/zero of=/tmp/testfile bs=1M count=2 && sync", 30*time.Second)).ToNot(HaveOccurred())

	By("Running the requested operations and ensuring CBT is not curropted")
	switch op {
	case "migrate":
		By("Performing live migration")
		migration := libmigration.New(vmi.Name, vmi.Namespace)
		migration = libmigration.RunMigrationAndExpectToCompleteWithDefaultTimeout(virtClient, migration)
		libmigration.ConfirmVMIPostMigration(virtClient, vmi, migration)
	case "restart":
		By("Restarting VM")
		libvmops.StopVirtualMachine(vm)
		libvmops.StartVirtualMachine(vm)
	}
	cbtOverlayPath := fmt.Sprintf("/cbt/%s.qcow2", volumeName)
	checkCBTIntegrityForPersistenceTest(virtClient, vm, cbtOverlayPath)
}

func checkCBTIntegrityForPersistenceTest(virtClient kubecli.KubevirtClient, vm *v1.VirtualMachine, cbtOverlayPath string) {
	By("Checking CBT data integrity")
	vmi, err := virtClient.VirtualMachineInstance(vm.Namespace).Get(context.Background(), vm.Name, metav1.GetOptions{})
	Expect(err).ShouldNot(HaveOccurred())
	pvcName := backendstorage.CurrentPVCName(vmi)
	Expect(pvcName).ToNot(BeEmpty(), "Backend storage PVC name should not be empty")
	vm = libvmops.StopVirtualMachine(vm)
	pvc, err := virtClient.CoreV1().PersistentVolumeClaims(vm.Namespace).Get(context.Background(), pvcName, metav1.GetOptions{})
	Expect(err).ToNot(HaveOccurred())
	pod := libstorage.RenderPodWithPVC(
		"cbt-consumer-"+rand.String(5),
		[]string{"/bin/bash", "-c", "touch /tmp/startup; while true; do echo hello; sleep 2; done"},
		nil, pvc,
	)
	pod, err = libpod.Run(pod, vm.Namespace)
	ExpectWithOffset(1, err).ToNot(HaveOccurred())
	containerName := pod.Spec.Containers[0].Name
	stdout, err := exec.ExecuteCommandOnPod(pod, containerName, []string{"/bin/sh", "-c", fmt.Sprintf("qemu-img info %s/%s", libstorage.DefaultPvcMountPath, cbtOverlayPath)})
	Expect(err).ToNot(HaveOccurred())
	By("Verifying qcow2 file is not corrupted")
	Expect(stdout).To(ContainSubstring("corrupt: false"))
	err = virtClient.CoreV1().Pods(pod.Namespace).Delete(context.Background(), pod.Name, metav1.DeleteOptions{})
	Expect(err).ToNot(HaveOccurred())
}
