#************************************************************************************
# 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_power {
    echo '<1>  Reboot'
    echo '<2>  Halt'
    echo '<0>  Return to menu'
    echo -e '\nPlease enter your choice:'
    
    unset vtOpt
    read vtOpt
    
    if [ "$vtOpt" = "1" ]; then
        echo -e '\n\nSystem is rebooting ... \n'
        sleep 1
        reboot
    elif [ "$vtOpt" = "2" ]; then
        echo -e '\n\nSystem is halting ... \n'
        sleep 1
        halt
    fi
}

function get_os_type {
    set vtoy_os=Linux    
    for file in "efi/microsoft" "sources/boot.wim" "boot/bcd" "bootmgr.efi" "boot/etfsboot.com" "BOOT/etfsboot.com"; do
        if [ -e $1/$file ]; then
            set vtoy_os=Windows
            break
        fi
    done
    
    if [ -n "${vtdebug_flag}" ]; then
        echo ISO is $vtoy_os
    fi
}

function vt_check_pe {
    unset VT_PE_SUPPORT

    if [ -f $1/HBCD_PE.ini ]; then
        set ventoy_compatible=YES
        set VT_PE_SUPPORT=YES
    elif [ -f $1/easyu.flg ]; then
        set VT_PE_SUPPORT=YES
    elif [ -f $1/USM.ICO ]; then
        set VT_PE_SUPPORT=YES
    elif [ -d $1/USM_TOOL ]; then
        set VT_PE_SUPPORT=YES
    fi
}

function locate_initrd {
    vt_linux_locate_initrd 

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

function find_wim_file {
    unset ventoy_wim_file
    
    for file in "sources/boot.wim" "sources/BOOT.WIM" "Sources/Win10PEx64.WIM" "boot/BOOT.WIM" \
                "winpe_x64.wim" "boot/10pex64.wim" "BOOT/USM1PE6L.WIM" "BOOT/USM1PE6F.WIM"; do
        if [ -e $1/$file ]; then
            set ventoy_wim_file=$1/$file
            break
        fi
    done
}

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
        
    fi
}

function uefi_windows_menu_func {
    vt_windows_reset
    
    if [ "$ventoy_compatible" = "NO" ]; then        
        find_wim_file (loop)
        if [ -n "$ventoy_wim_file" ]; then
            vt_windows_locate_wim $ventoy_wim_file
        fi
    fi
    
    vt_windows_chain_data ${1}${chosen_path}

    if [ -n "${vtdebug_flag}" ]; then
        sleep 5
    fi

    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"
        sleep 5
    fi
}

function uefi_linux_menu_func {
    if [ "$ventoy_compatible" = "NO" ]; then    
        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
            set LoadIsoEfiDriver=on
            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
        
        if [ -e (loop)/syslinux/alt0/full.cz ]; then
            set LoadIsoEfiDriver=on
            set FirstTryBootFile='@EFI@BOOT@grubx64.efi'
        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
    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"
        sleep 5
    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

    if vt_is_udf ${1}${chosen_path}; then
        set ventoy_fs_probe=udf
    else
        set ventoy_fs_probe=iso9660
    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_pe (loop)        
        if [ "$VT_PE_SUPPORT" != "YES" ]; then
            if [ "$ventoy_fs_probe" = "iso9660" ]; then        
                set ventoy_compatible=YES
            fi
        fi
    
        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                
        find_wim_file (loop)
        if [ -n "$ventoy_wim_file" ]; then
            vt_windows_locate_wim $ventoy_wim_file
        elif [ -n "${vtdebug_flag}" ]; then
            echo No wim file found
        fi
    fi
       
    vt_windows_chain_data ${1}${chosen_path}
    
    if [ -n "${vtdebug_flag}" ]; then
        sleep 5
    fi
    
    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"
        sleep 5
    fi
}

function legacy_linux_menu_func {
    if [ "$ventoy_compatible" = "NO" ]; then
        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}
    
    if [ -n "${vtdebug_flag}" ]; then
        sleep 5
    fi
    
    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"
        sleep 5
    fi
}

function legacy_iso_menu_func {

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

    vt_chosen_img_path chosen_path

    if vt_is_udf ${1}${chosen_path}; then
        set ventoy_fs_probe=udf
    else
        set ventoy_fs_probe=iso9660
    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_pe (loop)        
        if [ "$VT_PE_SUPPORT" != "YES" ]; then
            if [ "$ventoy_fs_probe" = "iso9660" ]; then        
                set ventoy_compatible=YES
            fi
        fi
        
        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 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      
}


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

set VENTOY_VERSION="1.0.11"

# 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="ISO RAW"
set VTOY_ISO_UEFI_DRV_STR="UEFI FS"

set VTOY_F2_CMD="ventoy_power"

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)
            break
        fi
    done
else
    set vtoy_path=($root)/ventoy    
    set iso_path=($vtoy_dev,1)
fi

loadfont ascii

#Load Plugin
if [ -f $iso_path/ventoy/ventoy.json ]; then
   vt_load_plugin $iso_path
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"
else
    set VTOY_F3_CMD="vt_dynamic_menu 1 0"
    set VTOY_HOTKEY_TIP="F1:Memdisk  F2:Power  F3:ListView"
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

terminal_output  gfxterm

#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
    menuentry "No ISO files found (Press enter to reboot ...)" {
        echo -e "\n    Rebooting ... "
        reboot
    }
fi