from enum import IntEnum
from uuid import UUID

from virt.firmware.varstore import autodetect
from virt.firmware.efi import devpath
from virt.firmware.efi import efivar

class DevicePathType(IntEnum):
    MediaDevicePath = 4

class MediaDevicePathSubType(IntEnum):
    PIWGFirmwareFile = 6
    PIWGFirmwareVolume = 7

FvNameGuid_By_Type = {
    'ArmVirt': "64074afe-340a-4be6-94ba-91b5b4d0f71e",
    'Ovmf': '7cb8bdc9-f8eb-4f34-aaea-3ee4af6516a1',
    'LoongArchVirt': "5d19a5b3-130f-459b-a292-9270a9e6bc62"
}
UefiShellFileGuid = "7c04a583-9e3e-4f1c-ad65-e05268d0b4d1"

def init_shell_entry(varfile_name, varfile_type):
    fv_name_guid = FvNameGuid_By_Type[varfile_type]
    # After booting with an unitialized firmware volume, the shell ends up
    # being assigned to Boot0002. Let's stick with that.
    # Note: Boot0000 always seems to get overwritten to be the
    # BootManagerMenuApp on first boot.
    shell_index = 2

    varstore = autodetect.open_varstore(varfile_name)
    bpath = devpath.DevicePath()
    elem = devpath.DevicePathElem()
    elem.devtype = DevicePathType.MediaDevicePath
    elem.subtype = MediaDevicePathSubType.PIWGFirmwareVolume
    elem.data = UUID(fv_name_guid).bytes_le
    bpath.append(elem)

    elem = devpath.DevicePathElem()
    elem.devtype = DevicePathType.MediaDevicePath
    elem.subtype = MediaDevicePathSubType.PIWGFirmwareFile
    elem.data = UUID(UefiShellFileGuid).bytes_le
    bpath.append(elem)

    varlist = efivar.EfiVarList()
    varlist.set_boot_entry(
        shell_index,
        title="EFI Internal Shell",
        path=bpath,
    )
    varlist.set_boot_next(shell_index)
    varstore.write_varstore(varfile_name, varlist)