#************************************************************************************
# Copyright (c) 2020, longpanda <admin@ventoy.net>
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.
# 
#************************************************************************************

function ventoy_pause {
    if [ -n "${vtdebug_flag}" ]; then
        echo "press Enter to continue ......"
        read vtTmpPause
    fi
}

function ventoy_debug_pause {
    if [ -n "${vtdebug_flag}" ]; then
        echo "press Enter to continue ......"
        read vtTmpPause
    fi
}


function ventoy_power {
    configfile $prefix/power.cfg
}

function ventoy_diagnosis {    
    configfile $prefix/debug.cfg
}

function ventoy_localboot {    
    configfile $prefix/localboot.cfg
}

function get_os_type {
    set vtoy_os=Linux
    
    for file in "efi/microsoft/boot/bcd" "sources/boot.wim" "boot/bcd" "bootmgr.efi" "boot/etfsboot.com" ; do        
        if vt_file_exist_nocase (loop)/$file; then        
            set vtoy_os=Windows
            break
        fi
    done

    if [ -n "${vtdebug_flag}" ]; then
        echo ISO is $vtoy_os
    fi
}

function vt_check_compatible_pe {
    #Check for PE without external tools
    if [ -f $1/HBCD_PE.ini ]; then
        set ventoy_compatible=YES        
    fi
}

function locate_initrd {
    vt_linux_locate_initrd 

    if [ -n "${vtdebug_flag}" ]; then        
        vt_linux_dump_initrd
        ventoy_debug_pause
    fi
}

function locate_wim {
    vt_windows_locate_wim_patch (loop)
    
    if [ -n "${vtdebug_flag}" ]; then
        echo '###############################################'
        vt_dump_wim_patch
        echo '###############################################'
        ventoy_debug_pause
    fi
}

function distro_specify_wim_patch {
    if [ -d (loop)/h3pe ]; then
        vt_windows_collect_wim_patch wim /BOOT/H3_10PE.WIM
        vt_windows_collect_wim_patch wim /BOOT/H3_7PE.WIM
        vt_windows_collect_wim_patch wim /BOOT/H3_8PE.WIM
        vt_windows_collect_wim_patch wim /BOOT/H3_81PE.WIM
    fi
}

function distro_specify_wim_patch_phase2 {
    if [ -f (loop)/boot/boot.wim ]; then
        vt_windows_collect_wim_patch wim /boot/boot.wim
    fi
}


function distro_specify_initrd_file {
    if [ -e (loop)/boot/all.rdz ]; then
        vt_linux_specify_initrd_file /boot/all.rdz
    elif [ -e (loop)/boot/xen.gz ]; then 
        if [ -e (loop)/install.img ]; then
            vt_linux_specify_initrd_file /install.img
        fi
    elif [ -d (loop)/casper ]; then 
        if [ -e (loop)/casper/initrd ]; then
            vt_linux_specify_initrd_file /casper/initrd
        fi
        if [ -e (loop)/casper/initrd-oem ]; then
            vt_linux_specify_initrd_file /casper/initrd-oem
        fi
    elif [ -e (loop)/boot/grub/initrd.xz ]; then
        vt_linux_specify_initrd_file /boot/grub/initrd.xz
    elif [ -e (loop)/initrd.gz ]; then
        vt_linux_specify_initrd_file /initrd.gz
    elif [ -e (loop)/slax/boot/initrfs.img ]; then
        vt_linux_specify_initrd_file /slax/boot/initrfs.img
    elif [ -e (loop)/pmagic/initrd.img ]; then
        vt_linux_specify_initrd_file /pmagic/initrd.img
    elif [ -e (loop)/boot/initrd.xz ]; then
        vt_linux_specify_initrd_file /boot/initrd.xz
    elif [ -f (loop)/boot/initrd ]; then
        vt_linux_specify_initrd_file /boot/initrd
    elif [ -f (loop)/boot/x86_64/loader/initrd ]; then
        vt_linux_specify_initrd_file /boot/x86_64/loader/initrd
    elif [ -f (loop)/boot/initramfs-x86_64.img ]; then
        vt_linux_specify_initrd_file /boot/initramfs-x86_64.img
        

    fi
}


function distro_specify_initrd_file_phase2 {
    if [ -f (loop)/boot/initrd.img ]; then
        vt_linux_specify_initrd_file /boot/initrd.img
    elif [ -f (loop)/Setup/initrd.gz ]; then
        vt_linux_specify_initrd_file /Setup/initrd.gz
    elif [ -f (loop)/isolinux/initramfs ]; then
        vt_linux_specify_initrd_file /isolinux/initramfs
    elif [ -f (loop)/boot/iniramfs.igz ]; then
        vt_linux_specify_initrd_file /boot/iniramfs.igz
    elif [ -f (loop)/initrd-x86_64 ]; then
        vt_linux_specify_initrd_file /initrd-x86_64
    elif [ -f (loop)/live/initrd.img ]; then 
        vt_linux_specify_initrd_file /live/initrd.img
    elif [ -f (loop)/initrd.img ]; then 
        vt_linux_specify_initrd_file /initrd.img

    fi
}

function uefi_windows_menu_func {
    vt_windows_reset

    if [ "$ventoy_compatible" = "NO" ]; then   
        
        if [ "$ventoy_fs_probe" = "iso9660" ]; then
            loopback -d loop
            vt_iso9660_nojoliet 1            
            loopback loop $1$2
        fi
        
        for file in "efi/microsoft/boot/bcd"; do
            vt_windows_collect_wim_patch bcd (loop)/$file                
        done

        vt_windows_count_wim_patch vt_wim_cnt
        if [ $vt_wim_cnt -eq 0 ]; then
            distro_specify_wim_patch_phase2
        fi
        
        ventoy_debug_pause        
        locate_wim
    fi

    vt_windows_chain_data ${1}${chosen_path}    
    ventoy_debug_pause
    
    if [ -n "$vtoy_chain_mem_addr" ]; then
        terminal_output  console       
        chainloader ${vtoy_path}/ventoy_x64.efi  env_param=${env_param} isoefi=${LoadIsoEfiDriver} ${vtdebug_flag} mem:${vtoy_chain_mem_addr}:size:${vtoy_chain_mem_size}
        boot
    else
        echo "chain empty failed"
        ventoy_pause
    fi
}

function uefi_linux_menu_func {
    if [ "$ventoy_compatible" = "NO" ]; then    
        
        if [ "$ventoy_fs_probe" = "udf" ]; then
            loopback -d loop            
            set ventoy_fs_probe=iso9660
            loopback loop $1$2
        fi
        
        vt_load_cpio  ${vtoy_path}/ventoy.cpio   $2 $1

        vt_linux_clear_initrd
        
        if [ -d (loop)/pmagic ]; then
            vt_linux_specify_initrd_file /pmagic/initrd.img
        else
            for file in "boot/grub/grub.cfg" "EFI/BOOT/grub.cfg" "EFI/boot/grub.cfg" "efi/boot/grub.cfg" "EFI/BOOT/BOOTX64.conf" "/grub/grub.cfg"; do
                if [ -e (loop)/$file ]; then                    
                    vt_linux_parse_initrd_grub  file  (loop)/$file
                fi
            done
        fi

        # special process for special distros
        if [ -d (loop)/loader/entries ]; then
            vt_linux_parse_initrd_grub  dir  (loop)/loader/entries/
        elif [ -d (loop)/boot/grub ]; then
            vt_linux_parse_initrd_grub  dir  (loop)/boot/grub/
        fi
        
        distro_specify_initrd_file
        
        vt_linux_initrd_count vtcount
        if [ $vtcount -eq 0 ]; then
            distro_specify_initrd_file_phase2
            
            if [ "$vt_efi_dir" = "NO" ]; then
                if [ -f (loop)/efi.img ];  then
                    vt_add_replace_file 0 "initrd"
                fi
            fi
        fi
        
        locate_initrd
        
        if [ -d (loop)/loader/entries ]; then
            vt_linux_get_main_initrd_index vtindex
            
            if [ -d (loop)/arch ]; then
                if [ -f (loop)/arch/boot/x86_64/archiso.img ]; then
                    vt_add_replace_file $vtindex "EFI\\archiso\\archiso.img"
                elif [ -f (loop)/boot/initramfs_x86_64.img ]; then
                    vt_add_replace_file $vtindex "boot\\initramfs_x86_64.img"
                fi
            elif [ -f (loop)/EFI/BOOT/initrd.gz ]; then
                vt_add_replace_file $vtindex "EFI\\BOOT\\initrd.gz"
            fi
        elif [ -e (loop)/syslinux/alt0/full.cz ]; then
            vt_add_replace_file 0 "EFI\\BOOT\\full.cz"            
            set FirstTryBootFile='@EFI@BOOT@grubx64.efi'
        fi
        
    fi
    
    vt_linux_chain_data ${1}${chosen_path}

    if [ -n "$vtoy_chain_mem_addr" ]; then
        terminal_output  console       
        chainloader ${vtoy_path}/ventoy_x64.efi  env_param=${env_param} isoefi=${LoadIsoEfiDriver} FirstTry=${FirstTryBootFile} ${vtdebug_flag} mem:${vtoy_chain_mem_addr}:size:${vtoy_chain_mem_size}
        boot
    else
        echo "chain empty failed"
        ventoy_pause
    fi
}


function uefi_iso_menu_func {

    if [ -d (loop)/ ]; then
        loopback -d loop
    fi

    if [ -n "$vtisouefi" ]; then
        set LoadIsoEfiDriver=on
        unset vtisouefi
    elif vt_check_mode 2; then
        set LoadIsoEfiDriver=on
    else
        unset LoadIsoEfiDriver
    fi

    vt_chosen_img_path chosen_path
    vt_select_auto_install ${chosen_path}
    vt_select_persistence ${chosen_path}

    if vt_is_udf ${1}${chosen_path}; then
        set ventoy_fs_probe=udf
    else
        set ventoy_fs_probe=iso9660
        vt_iso9660_nojoliet 0
    fi

    loopback loop ${1}${chosen_path}
    get_os_type (loop)
    
    if [ -d (loop)/EFI ]; then
        set vt_efi_dir=YES
    elif [ -d (loop)/efi ]; then
        set vt_efi_dir=YES
    else
        set vt_efi_dir=NO
    fi
    
    if [ -n "$vtcompat" ]; then
        set ventoy_compatible=YES
        unset vtcompat
    elif vt_check_mode 1; then
        set ventoy_compatible=YES
    else
        vt_check_compatible (loop)
    fi
    
    vt_img_sector ${1}${chosen_path}
    
    if [ "$vtoy_os" = "Windows" ]; then
        vt_check_compatible_pe (loop)        
        uefi_windows_menu_func  $1 ${chosen_path}
    else
        uefi_linux_menu_func  $1 ${chosen_path}
    fi

    terminal_output  gfxterm    
}

function uefi_iso_memdisk {    
    vt_chosen_img_path chosen_path

    echo 'Loading ISO file to memory ...'
    vt_load_iso_to_mem ${1}${chosen_path} vtoy_iso_buf
    
    terminal_output  console
    chainloader ${vtoy_path}/ventoy_x64.efi memdisk env_param=${env_param} isoefi=${LoadIsoEfiDriver} ${vtdebug_flag} mem:${vtoy_iso_buf_addr}:size:${vtoy_iso_buf_size}
    boot
}


function legacy_windows_menu_func {
    vt_windows_reset
    
    if [ "$ventoy_compatible" = "NO" ]; then   
        
        if [ "$ventoy_fs_probe" = "iso9660" ]; then
            loopback -d loop
            vt_iso9660_nojoliet 1
            loopback loop $1$2
        fi
        
        for file in "boot/bcd" "/efi/microsoft/boot/bcd" "SSTR/BCD"; do
            vt_windows_collect_wim_patch bcd (loop)/$file                
        done
        
        distro_specify_wim_patch

        vt_windows_count_wim_patch vt_wim_cnt
        if [ $vt_wim_cnt -eq 0 ]; then
            distro_specify_wim_patch_phase2
        fi
        
        ventoy_debug_pause        
        locate_wim
    fi

    vt_windows_chain_data ${1}${chosen_path}    
    ventoy_debug_pause    
    
    if [ -n "$vtoy_chain_mem_addr" ]; then
        linux16   $vtoy_path/ipxe.krn ${vtdebug_flag} ibft mem:${vtoy_chain_mem_addr}:size:${vtoy_chain_mem_size}
        boot
    else
        echo "chain empty failed"
        ventoy_pause
    fi
}

function legacy_linux_menu_func {
    if [ "$ventoy_compatible" = "NO" ]; then

        if [ "$ventoy_fs_probe" = "udf" ]; then
            loopback -d loop
            set ventoy_fs_probe=iso9660
            loopback loop $1$2
        fi
    
        vt_load_cpio  $vtoy_path/ventoy.cpio  $2 $1

        vt_linux_clear_initrd
        
        if [ -d (loop)/pmagic ]; then
            vt_linux_specify_initrd_file /pmagic/initrd.img
        else
            for dir in "isolinux" "boot/isolinux" "boot/x86_64/loader" "syslinux" "boot/syslinux"; do
                if [ -d (loop)/$dir ]; then
                    vt_linux_parse_initrd_isolinux   (loop)/$dir/
                fi
            done
        fi
        
        # special process for special distros
        #archlinux
        if [ -d (loop)/arch/boot/syslinux ]; then
            vt_linux_parse_initrd_isolinux   (loop)/arch/boot/syslinux/  /arch/
            vt_linux_parse_initrd_isolinux   (loop)/arch/boot/syslinux/  /arch/boot/syslinux/
            
        #manjaro
        elif [ -d (loop)/manjaro ]; then
            if [ -e (loop)/boot/grub/kernels.cfg ]; then
                vt_linux_parse_initrd_grub  file  (loop)/boot/grub/kernels.cfg
            fi
        elif [ -e (loop)/boot/grub/grub.cfg ]; then                
            vt_linux_parse_initrd_grub  file  (loop)/boot/grub/grub.cfg
        fi
        
        distro_specify_initrd_file
        
        vt_linux_initrd_count vtcount
        if [ $vtcount -eq 0 ]; then
            distro_specify_initrd_file_phase2
        fi
        
        locate_initrd
    fi
    
    vt_linux_chain_data ${1}${chosen_path}    
    ventoy_debug_pause
    
    if [ -n "$vtoy_chain_mem_addr" ]; then
        linux16   $vtoy_path/ipxe.krn ${vtdebug_flag}  mem:${vtoy_chain_mem_addr}:size:${vtoy_chain_mem_size}
        boot
    else
        echo "chain empty failed"
        ventoy_pause
    fi
}

function legacy_iso_menu_func {

    if [ -d (loop)/ ]; then
        loopback -d loop
    fi

    vt_chosen_img_path chosen_path
    vt_select_auto_install ${chosen_path}
    vt_select_persistence ${chosen_path}

    if vt_is_udf ${1}${chosen_path}; then
        set ventoy_fs_probe=udf
    else
        set ventoy_fs_probe=iso9660
        vt_iso9660_nojoliet 0
    fi
    
    loopback loop ${1}${chosen_path}
    get_os_type (loop)

    if [ -n "$vtcompat" ]; then
        set ventoy_compatible=YES
        unset vtcompat
    elif vt_check_mode 1; then
        set ventoy_compatible=YES
    else
        vt_check_compatible (loop)
    fi
    
    vt_img_sector ${1}${chosen_path}

    if [ "$vtoy_os" = "Windows" ]; then
        vt_check_compatible_pe (loop)        
        legacy_windows_menu_func  $1 ${chosen_path}
    else
        legacy_linux_menu_func  $1 ${chosen_path}
    fi
}

function legacy_iso_memdisk {
    vt_chosen_img_path chosen_path

    linux16   $vtoy_path/memdisk iso raw    
    echo "Loading ISO file to memory ..."
    initrd16  ${1}${chosen_path}
    boot
}

function iso_common_menuentry {
    if [ "$grub_platform" = "pc" ]; then
        if vt_check_mode 0; then
            legacy_iso_memdisk $iso_path
        else
            legacy_iso_menu_func $iso_path
        fi
    else
        if vt_check_mode 0; then
            uefi_iso_memdisk $iso_path 
        else
            uefi_iso_menu_func $iso_path                
        fi
    fi      
}

function iso_unsupport_menuentry {
    echo -e "\n The name of the iso file could NOT contain space or non-ascii characters. \n"
    echo -e "\n Will return to main menu after 10 seconds ...\n"
    sleep 10
}


function wim_common_menuentry {
    vt_chosen_img_path chosen_path
    vt_wim_chain_data ${iso_path}${chosen_path}
    
    ventoy_debug_pause    
    
    if [ -n "$vtoy_chain_mem_addr" ]; then
        if [ "$grub_platform" = "pc" ]; then
            linux16   $vtoy_path/ipxe.krn ${vtdebug_flag}  mem:${vtoy_chain_mem_addr}:size:${vtoy_chain_mem_size}
        else
            terminal_output  console       
            chainloader ${vtoy_path}/ventoy_x64.efi  env_param=${env_param} isoefi=${LoadIsoEfiDriver} ${vtdebug_flag} mem:${vtoy_chain_mem_addr}:size:${vtoy_chain_mem_size}
        fi
        boot
    else
        echo "chain empty failed"
        ventoy_pause
    fi
}

function wim_unsupport_menuentry {
    echo -e "\n The name of the wim file could NOT contain space or non-ascii characters. \n"
    echo -e "\n Will return to main menu after 10 seconds ...\n"
    sleep 10
}

#############################################################
#############################################################
#############################################################
#######                 Main Process              ###########
#############################################################
#############################################################
#############################################################

set VENTOY_VERSION="1.0.13"

# Default menu display mode, you can change it as you want.
#    0: List mode   
#    1: TreeView mode
set VTOY_DEFAULT_MENU_MODE=0

#disable timeout
unset timeout

set VTOY_MEM_DISK_STR="[Memdisk]"
set VTOY_ISO_RAW_STR="Compatible Mode"
set VTOY_ISO_UEFI_DRV_STR="UEFI FS"

set VTOY_F2_CMD="ventoy_power"
set VTOY_F4_CMD="ventoy_localboot"
set VTOY_F5_CMD="ventoy_diagnosis"

if [ "$grub_platform" = "pc" ]; then
    set VTOY_TEXT_MENU_VER="Ventoy $VENTOY_VERSION BIOS  www.ventoy.net"
else
    set VTOY_TEXT_MENU_VER="Ventoy $VENTOY_VERSION UEFI  www.ventoy.net"
fi

vt_device $root  vtoy_dev

if [ "$vtoy_dev" = "tftp" ]; then
    set vtoy_path=($root)
    for vtid in 0 1 2 3; do
        if [ -d (hd$vtid,2)/ventoy ]; then
            set iso_path=(hd$vtid,1)
            set vtoy_efi_part=(hd$vtid,2)
            break
        fi
    done
    loadfont ascii
else
    if [ "$prefix" = "(ventoydisk)/grub" ]; then
        set vtoy_path=(ventoydisk)/ventoy
    else
        set vtoy_path=($root)/ventoy
    fi

    set iso_path=($vtoy_dev,1)
    set vtoy_efi_part=($vtoy_dev,2)
    loadfont unicode
fi


#Load Plugin
if [ -f $iso_path/ventoy/ventoy.json ]; then
   vt_load_plugin $iso_path
fi

if [ -f $iso_path/ventoy/ventoy_wimboot.img ]; then
    vt_load_wimboot $iso_path/ventoy/ventoy_wimboot.img
elif [ -f $vtoy_efi_part/ventoy/ventoy_wimboot.img ]; then
    vt_load_wimboot $vtoy_efi_part/ventoy/ventoy_wimboot.img
fi


if [ $VTOY_DEFAULT_MENU_MODE -eq 0 ]; then
    set VTOY_F3_CMD="vt_dynamic_menu 1 1"
    set VTOY_HOTKEY_TIP="F1:Memdisk  F2:Power  F3:TreeView  F4:Localboot  F5:Debug"
else
    set VTOY_F3_CMD="vt_dynamic_menu 1 0"
    set VTOY_HOTKEY_TIP="F1:Memdisk  F2:Power  F3:ListView  F4:Localboot  F5:Debug"
fi


if [ -n "$vtoy_gfxmode" ]; then
    set gfxmode=$vtoy_gfxmode
else
    set gfxmode=1920x1080,1366x768,1024x768
fi

if [ -n "$vtoy_theme" ]; then
    set theme=$vtoy_theme
else
    set theme=$prefix/themes/ventoy/theme.txt
fi

if [ "$vtoy_display_mode" = "CLI" ]; then
    terminal_output  console
else
    terminal_output  gfxterm
fi

#vtdebug on

#colect all image files (iso files)
set ventoy_img_count=0
vt_list_img $iso_path ventoy_img_count

#Main menu
if [ $ventoy_img_count -gt 0 ]; then
    if [ $VTOY_DEFAULT_MENU_MODE -eq 0 ]; then
        vt_dynamic_menu 0 0 
    else
        vt_dynamic_menu 0 1
    fi
else
    if [ -n "$VTOY_NO_ISO_TIP" ]; then
        NO_ISO_MENU="No ISO files found, $VTOY_NO_ISO_TIP"
    elif [ -n "$VTOY_DEFAULT_SEARCH_ROOT" ]; then
        NO_ISO_MENU="No ISO files found, please check VTOY_DEFAULT_SEARCH_ROOT"
    else
        NO_ISO_MENU="No ISO files found"
    fi
    menuentry "$NO_ISO_MENU (Press enter to reboot ...)" {
        echo -e "\n    Rebooting ... "
        reboot
    }
fi
