#!/usr/bin/env python3
"""
Patch Ventoy 1.1.18 grub.cfg for the custom Xorboot layout:

  * data (ISO) partition is detected by marker file /ISO/Ventoy/NOT.DEL
    instead of being hardcoded to partition 1;
  * partition numbers are exported to the ventoy GRUB module
    (VTOY_DATA_PART_IDX / VTOY_EFI_PART_IDX);
  * user config + custom menu live in /ISO/Ventoy/ instead of /ventoy/.

Usage:
    python3 patch_ventoy_grubcfg.py <path-to-vtoyefi/grub/grub.cfg>
"""

import io
import sys

MARK = "# VTOY-CUSTOM-LAYOUT"

OLD_BLOCK = (
    "    set vtoydev=$vtoy_dev\n"
    "    set vtoy_iso_part=($vtoy_dev,1)\n"
    "    set vtoy_efi_part=($vtoy_dev,2)\n"
    "    \n"
    "    vt_load_file_to_mem \"auto\" $prefix/fonts/unicode.pf2 vtoy_font_mem\n"
)

NEW_BLOCK = (
    "    set vtoydev=$vtoy_dev\n"
    "    set vtoy_iso_part=($vtoy_dev,1)\n"
    "    set VTOY_DATA_PART_IDX=1\n"
    "    set vtoy_efi_part=($vtoy_dev,2)\n"
    "    set VTOY_EFI_PART_IDX=2\n"
    "\n"
    "    " + MARK + ": data (ISO) partition is the one that contains\n"
    "    # /ISO/Ventoy/NOT.DEL (fallback: partition 1). VTOYEFI stays partition 2.\n"
    "    for vt_part in 1 2 3 4; do\n"
    "        if [ -f ($vtoy_dev,$vt_part)/ISO/Ventoy/NOT.DEL ]; then\n"
    "            set vtoy_iso_part=($vtoy_dev,$vt_part)\n"
    "            set VTOY_DATA_PART_IDX=$vt_part\n"
    "            break\n"
    "        fi\n"
    "    done\n"
    "    export VTOY_DATA_PART_IDX\n"
    "    export VTOY_EFI_PART_IDX\n"
    "\n"
    "    vt_load_file_to_mem \"auto\" $prefix/fonts/unicode.pf2 vtoy_font_mem\n"
)

REPLACEMENTS = [
    # plugin json at /ISO/Ventoy/ (both tftp and normal branches)
    ("$vtoy_iso_part/ventoy/ventoy.json", "$vtoy_iso_part/ISO/Ventoy/ventoy.json"),
    # custom menu file at /ISO/Ventoy/
    ("$vt_plugin_path/ventoy/ventoy_grub.cfg", "$vt_plugin_path/ISO/Ventoy/ventoy_grub.cfg"),
    # default menu language: Russian (Ventoy opens in Russian unless user picks)
    ('    vt_init_menu_lang en_US\n', '    vt_init_menu_lang ru_RU\n'),
    # WinPE with a non-standard BCD (Sergei Strelec: SSTR/BCD, WePE: WEPE/B64...).
    # Ventoy reads the list of wims to patch from BCD files; the x86 wims of such
    # PEs are only referenced there, so without this the mount under ia32 never
    # happens. Same set as the legacy branch uses.
    ('for file in "efi/microsoft/boot/bcd"; do',
     'for file in "efi/microsoft/boot/bcd" "boot/bcd" "SSTR/BCD" "WEPE/B64" "WEPE/B32" "WEPE/BCD" "ISPE/BCD" "ISYL/BCD" "2K10/FONTS/BCD" "2K10/BCD" "boot/uqi/BCD"; do'),
    ('for file in "boot/bcd" "/efi/microsoft/boot/bcd" "SSTR/BCD" "boot/bce"; do',
     'for file in "boot/bcd" "/efi/microsoft/boot/bcd" "SSTR/BCD" "WEPE/B64" "WEPE/B32" "WEPE/BCD" "ISPE/BCD" "ISYL/BCD" "2K10/FONTS/BCD" "2K10/BCD" "boot/uqi/BCD" "boot/bce"; do'),
]


def main():
    path = sys.argv[1]
    src = io.open(path, encoding="utf-8", newline="").read()

    if MARK in src:
        print("already patched:", path)
        return

    if src.count(OLD_BLOCK) != 1:
        raise SystemExit("layout block not found exactly once")

    src = src.replace(OLD_BLOCK, NEW_BLOCK, 1)
    print("layout block: marker-based data partition")

    for old, new in REPLACEMENTS:
        n = src.count(old)
        if n == 0:
            raise SystemExit("not found: %s" % old)
        src = src.replace(old, new)
        print("replaced %d x %s" % (n, old))

    io.open(path, "w", encoding="utf-8", newline="").write(src)
    print("written:", path)


if __name__ == "__main__":
    main()