#!/usr/bin/env python3
"""
Patch Ventoy 1.1.18 GRUB sources for the custom Xorboot layout:

  * data (ISO) partition is NOT partition 1 any more: it is selected at
    runtime by grub.cfg (marker file /ISO/Ventoy/NOT.DEL) and passed in
    VTOY_DATA_PART_IDX; VTOYEFI stays partition 2 (VTOY_EFI_PART_IDX).
  * the "official device" geometry check is disabled (it hardcodes
    data=part1 @1MB, VTOYEFI=part2 right after it).
  * plugin config is read from <data>/ISO/Ventoy/ventoy.json instead of
    <data>/ventoy/ventoy.json.

Run inside WSL against the build tree:
    python3 patch_ventoy_layout.py /root/vtbuild/GRUB2/SRC/grub-2.04/grub-core/ventoy
"""

import os
import re
import sys

MARK = "/* VTOY-CUSTOM-LAYOUT */"


def patch_cmd(path):
    src = open(path, encoding="utf-8", errors="surrogateescape").read()
    changed = False

    # --- 1. disable official-device geometry check ------------------------
    anchor = "    char *label = NULL;\n    struct grub_partition *partition;\n\n"
    early = (
        "    char *label = NULL;\n"
        "    struct grub_partition *partition;\n"
        "\n"
        "    " + MARK + "\n"
        "    /* Custom layout: ISO partition is found by marker file and VTOYEFI\n"
        "       is partition 2, so the original geometry checks do not apply. */\n"
        "    (void)workaround; (void)file; (void)offset; (void)devname;\n"
        "    (void)fs; (void)mbr; (void)disk; (void)dev2; (void)label;\n"
        "    (void)partition;\n"
        "    return ventoy_set_check_result(0, NULL);\n"
        "\n"
        "    if (dev->disk == NULL || dev->disk->partition == NULL)\n"
    )
    if MARK in src:
        print("  ventoy_cmd.c: device check already patched")
    else:
        if anchor not in src:
            raise SystemExit("anchor for device check not found")
        src = src.replace(anchor, early, 1)
        changed = True
        print("  ventoy_cmd.c: device check disabled")

    # --- 2. partition size lookup via env --------------------------------
    old_fn = re.search(
        r"grub_uint64_t ventoy_get_vtoy_partsize\(int part\)\n\{.*?\n\}\n",
        src, re.S)
    if not old_fn:
        raise SystemExit("ventoy_get_vtoy_partsize not found")

    new_fn = (
        "grub_uint64_t ventoy_get_vtoy_partsize(int part)\n"
        "{\n"
        "    grub_uint64_t sectors;\n"
        "    const char *env;\n"
        "    int idx = part;\n"
        "\n"
        "    " + MARK + "\n"
        "    /* part 0 = data (ISO) partition, part 1 = VTOYEFI.\n"
        "       Real partition numbers come from grub.cfg (1-based). */\n"
        "    env = ventoy_get_env(part == 0 ? \"VTOY_DATA_PART_IDX\" : \"VTOY_EFI_PART_IDX\");\n"
        "    if (env)\n"
        "    {\n"
        "        idx = (int)grub_strtol(env, NULL, 10) - 1;\n"
        "    }\n"
        "    if (idx < 0)\n"
        "    {\n"
        "        idx = 0;\n"
        "    }\n"
        "    if (idx > 3)\n"
        "    {\n"
        "        idx = 3;\n"
        "    }\n"
        "\n"
        "    if (grub_strncmp(g_ventoy_part_info->Head.Signature, \"EFI PART\", 8) == 0)\n"
        "    {\n"
        "        sectors = g_ventoy_part_info->PartTbl[idx].LastLBA + 1 - g_ventoy_part_info->PartTbl[idx].StartLBA;\n"
        "    }\n"
        "    else\n"
        "    {\n"
        "        sectors = g_ventoy_part_info->MBR.PartTbl[idx].SectorCount;\n"
        "    }\n"
        "\n"
        "    return sectors * 512;\n"
        "}\n"
    )
    if marker_in(old_fn.group(0)):
        print("  ventoy_cmd.c: partsize already patched")
    else:
        src = src[:old_fn.start()] + new_fn + src[old_fn.end():]
        changed = True
        print("  ventoy_cmd.c: partsize now uses VTOY_DATA_PART_IDX")

    if changed:
        open(path, "w", encoding="utf-8", errors="surrogateescape").write(src)
    return changed


def marker_in(text):
    return MARK in text


def patch_plugin(path):
    src = open(path, encoding="utf-8", errors="surrogateescape").read()
    if "ISO/Ventoy/ventoy.json" in src:
        print("  ventoy_plugin.c: json path already patched")
        return False

    n = src.count('"%s/ventoy/ventoy.json"')
    if n != 2:
        raise SystemExit("expected 2 json path occurrences, found %d" % n)
    src = src.replace('"%s/ventoy/ventoy.json"', '"%s/ISO/Ventoy/ventoy.json"')
    src = src.replace(
        'grub_printf("Plugin json file /ventoy/ventoy.json does NOT exist.\\n");',
        'grub_printf("Plugin json file /ISO/Ventoy/ventoy.json does NOT exist.\\n");')
    open(path, "w", encoding="utf-8", errors="surrogateescape").write(src)
    print("  ventoy_plugin.c: json path -> /ISO/Ventoy/ventoy.json")
    return True


def patch_vhd(path):
    """ventoy_vhd.c: VHD/IMG on the Ventoy data partition used partIndex = 0 (p1)."""
    src = open(path, encoding="utf-8", errors="surrogateescape").read()
    if MARK in src:
        print("  ventoy_vhd.c: already patched")
        return False

    pat = re.compile(r"(gpt = g_ventoy_part_info;[ \t]*\n)(\s*)partIndex = 0;")
    m = pat.search(src)
    if not m:
        raise SystemExit("vhd partIndex anchor not found")

    ind = m.group(2)
    new = (
        m.group(1)
        + ind + MARK + "\n"
        + ind + "/* data partition number comes from grub.cfg (VTOY_DATA_PART_IDX) */\n"
        + ind + "partIndex = 0;\n"
        + ind + "{\n"
        + ind + "    const char *vtEnv = ventoy_get_env(\"VTOY_DATA_PART_IDX\");\n"
        + ind + "    if (vtEnv)\n"
        + ind + "    {\n"
        + ind + "        partIndex = (int)grub_strtol(vtEnv, NULL, 10) - 1;\n"
        + ind + "    }\n"
        + ind + "    if (partIndex < 0)\n"
        + ind + "    {\n"
        + ind + "        partIndex = 0;\n"
        + ind + "    }\n"
        + ind + "    if (partIndex > 3)\n"
        + ind + "    {\n"
        + ind + "        partIndex = 3;\n"
        + ind + "    }\n"
        + ind + "}"
    )
    src = src[:m.start()] + new + src[m.end():]
    open(path, "w", encoding="utf-8", errors="surrogateescape").write(src)
    print("  ventoy_vhd.c: partIndex now uses VTOY_DATA_PART_IDX")
    return True


def main():
    d = sys.argv[1] if len(sys.argv) > 1 else "."
    patch_cmd(os.path.join(d, "ventoy_cmd.c"))
    patch_plugin(os.path.join(d, "ventoy_plugin.c"))
    patch_vhd(os.path.join(d, "ventoy_vhd.c"))
    print("OK")


if __name__ == "__main__":
    main()