#!/usr/bin/env python3
"""
Patch GRUB 2.12 source tree to compile the Ventoy module.
Fixes all 60 compilation errors by:
1. Adding missing struct fields to GRUB headers
2. Adding compat typedefs/macros/globals to ventoy_def.h
3. Providing working implementations for EFI allocation functions
4. Providing stubs for functions that need patched core files
5. Fixing grub_strtoul signature change
"""
import os, sys

SRC = sys.argv[1] if len(sys.argv) > 1 else "/mnt/data/Builds/hiperiso/build/grub2-212/SRC/grub-2.12"

def patch_file(path, old, new, must_match=True):
    full = os.path.join(SRC, path)
    with open(full) as f:
        c = f.read()
    if old not in c:
        if must_match:
            print(f"ERROR: pattern not found in {path}")
            sys.exit(1)
        else:
            print(f"SKIP: pattern not found in {path}")
            return
    c = c.replace(old, new, 1)
    with open(full, 'w') as f:
        f.write(c)
    print(f"PATCHED: {path}")

# ============================================================
# 1. Patch GRUB 2.12 headers: add missing struct fields
# ============================================================

# file.h: add int vlnk to struct grub_file
patch_file("include/grub/file.h",
    "  /* If file is not easily seekable. Should be set by underlying layer.  */\n  int not_easily_seekable;",
    "  /* If file is not easily seekable. Should be set by underlying layer.  */\n  int not_easily_seekable;\n\n  /* Ventoy vlnk flag.  */\n  int vlnk;")

# fs.h: add grub_uint64_t size to grub_dirhook_info
patch_file("include/grub/fs.h",
    "  grub_int64_t mtime;\n  grub_uint64_t inode;\n};",
    "  grub_int64_t mtime;\n  grub_uint64_t inode;\n  grub_uint64_t size;\n};")

# partition.h: add grub_uint64_t gpt_attrib to struct grub_partition
patch_file("include/grub/partition.h",
    "  grub_uint8_t msdostype;\n};",
    "  grub_uint8_t msdostype;\n\n  /* Ventoy: GPT attribute field for priority detection.  */\n  grub_uint64_t gpt_attrib;\n};")

# ============================================================
# 2. Fix grub_strtoul call in ventoy_json.c
# ============================================================
patch_file("grub-core/ventoy/ventoy_json.c",
    "Value = grub_strtoul(pcData, (char **)ppcEnd, 10);",
    "Value = grub_strtoul(pcData, (const char ** const)ppcEnd, 10);")

# ============================================================
# 2b. Patch kern/file.c: add mem: protocol support
# ============================================================
file_c_path = os.path.join(SRC, "grub-core/kern/file.c")
with open(file_c_path) as f:
    fc = f.read()

memfile_func = """
static grub_ssize_t
grub_memfile_read (grub_file_t file, char *buf, grub_size_t len)
{
  grub_size_t avail = file->size - file->offset;
  if (len > avail)
    len = avail;
  if (len == 0)
    return 0;
  grub_memcpy (buf, (const char *) file->data + file->offset, len);
  return (grub_ssize_t) len;
}

static grub_err_t
grub_memfile_close (grub_file_t file)
{
  return GRUB_ERR_NONE;
}

static struct grub_fs grub_mem_fs =
{
  .name = "mem",
  .fs_read = grub_memfile_read,
  .fs_close = grub_memfile_close,
};

static grub_file_t
grub_memfile_open (const char *name)
{
  grub_file_t file;
  const char *size_str;

  file = (grub_file_t) grub_zalloc (sizeof (*file));
  if (!file)
    return 0;

  file->name = grub_strdup (name);
  file->data = (void *) grub_strtoul (name + 4, NULL, 0);
  file->fs = &grub_mem_fs;

  size_str = grub_strstr (name, "size:");
  if (size_str)
    file->size = (grub_off_t) grub_strtoul (size_str + 5, NULL, 0);

  grub_errno = GRUB_ERR_NONE;
  return file;
}
"""
fc = fc.replace(
    "void (*EXPORT_VAR (grub_grubnet_fini)) (void);\n",
    "void (*EXPORT_VAR (grub_grubnet_fini)) (void);\n" + memfile_func, 1)

fc = fc.replace(
    "  grub_file_filter_id_t filter;\n\n  /* Reset grub_errno",
    "  grub_file_filter_id_t filter;\n\n  if (grub_strncmp (name, \"mem:\", 4) == 0)\n    return grub_memfile_open (name);\n\n  /* Reset grub_errno", 1)

with open(file_c_path, 'w') as f:
    f.write(fc)
print("PATCHED: grub-core/kern/file.c (mem: protocol)")

# ============================================================
# 2c. Patch font/font.c: add mem: to direct open path
# ============================================================
patch_file("grub-core/font/font.c",
    'if (filename[0] == \'(\' || filename[0] == \'/\' || filename[0] == \'+\')\n    file = grub_buffile_open (filename, GRUB_FILE_TYPE_FONT, 1024);',
    'if (filename[0] == \'(\' || filename[0] == \'/\' || filename[0] == \'+\'\n      || grub_strncmp (filename, "mem:", 4) == 0)\n    file = grub_buffile_open (filename, GRUB_FILE_TYPE_FONT, 1024);')

# ============================================================
# 3. Add missing includes to ventoy_def.h + append compat block
# ============================================================
ventoy_def_path = os.path.join(SRC, "grub-core/ventoy/ventoy_def.h")
with open(ventoy_def_path) as f:
    vd = f.read()

# ventoy_def.h has zero includes — relies on caller.
# Add them directly so all types resolve from any .c file.
include_anchor = "#define __VENTOY_DEF_H__\n"
if include_anchor in vd:
    vd = vd.replace(include_anchor,
        include_anchor + "\n#include <grub/command.h>\n#include <grub/extcmd.h>\n#include <grub/ventoy.h>\n", 1)
    print("PATCHED: grub-core/ventoy/ventoy_def.h (includes added)")
else:
    print("SKIP: ventoy_def.h include anchor not found")

COMPAT_BLOCK = """
/* ===== GRUB 2.12 compatibility layer for Ventoy module ===== */

/* Type: grub_efi_guid_t was renamed to grub_guid_t in GRUB 2.12 */
typedef grub_guid_t grub_efi_guid_t;

/* SB policy globals (from modsrc include/grub/env.h) */
#define VTOY_SB_POLICY_BYPASS   0
#define VTOY_SB_POLICY_CHECK    1
extern grub_uint8_t g_sys_sb;
extern grub_uint8_t g_sb_policy;

/* File type macro */
#define GRUB_FILE_TYPE_NO_VLNK GRUB_FILE_TYPE_NONE

/* Missing function declarations (from modsrc core patches) */
int ventoy_check_file_exist(const char *fmt, ...);
int grub_file_is_vlnk_suffix(const char *name, int len);
int grub_file_add_vlnk(const char *src, const char *dst);
int grub_file_vtoy_vlnk(const char *src, const char *dst);
const char *grub_file_get_vlnk(const char *name, int *vlnk);
grub_fs_t grub_fs_list_probe(grub_device_t device, const char **list);
void *grub_efi_allocate_iso_buf(grub_uint64_t size);
void *grub_efi_allocate_chain_buf(grub_uint64_t size);
void grub_efi_get_reserved_page_num(grub_uint64_t *total, grub_uint64_t *org_required, grub_uint64_t *new_required);
void ventoy_env_hook_root(int hook);
grub_err_t grub_register_vtoy_menu_lang_hook(grub_env_read_hook_t read_hook);
grub_err_t grub_disk_blocklist_read(void *chunklist, grub_uint64_t sector, grub_uint64_t size, grub_uint32_t log_sector_size);

/* grub_efi_get_variable compat: 2.12 changed to 4-arg returning status.
   Provide 3-arg wrapper returning void* (matching 2.04 API). */
#define grub_efi_get_variable ventoy_efi_get_variable_compat
void *ventoy_efi_get_variable_compat(const char *var, const grub_guid_t *guid, grub_size_t *size);

/* ===== End compatibility layer ===== */
"""

# Insert before the final #endif
endif_line = vd.rfind('#endif')
if endif_line < 0:
    print("ERROR: could not find #endif in ventoy_def.h")
    sys.exit(1)
vd = vd[:endif_line] + COMPAT_BLOCK + "\n" + vd[endif_line:]
with open(ventoy_def_path, 'w') as f:
    f.write(vd)
print("PATCHED: grub-core/ventoy/ventoy_def.h (compat block appended)")

# Disable Ventoy anti-tamper check (VTOY_CMD_CHECK calls grub_exit()
# when ESP partition isn't exactly 33554432 bytes — hiperiso layout
# may differ. Neutralize it.
VTOY_CHECK_OLD = "#define VTOY_CMD_CHECK(a) if (33554432 != g_ventoy_disk_part_size[a]) ventoy_syscall0(exit)"
VTOY_CHECK_NEW = "#define VTOY_CMD_CHECK(a) /* disabled for hiperiso */"
if VTOY_CHECK_OLD in vd:
    vd = vd.replace(VTOY_CHECK_OLD, VTOY_CHECK_NEW, 1)
    print("PATCHED: ventoy_def.h (VTOY_CMD_CHECK disabled)")
else:
    print("SKIP: VTOY_CMD_CHECK anchor not found (may already be patched)")
with open(ventoy_def_path, 'w') as f:
    f.write(vd)

# Bypass ventoy_check_official_device — it requires ventoy/ventoy.cpio
# on the ESP which hiperiso doesn't ship. Insert return 0 after the
# "cpio open failed" error path.
ventoy_c_path = os.path.join(SRC, "grub-core/ventoy/ventoy_cmd.c")
with open(ventoy_c_path) as f:
    vc = f.read()
cpio_marker = 'return ventoy_set_check_result(3 | 0x1000, "File ventoy/ventoy.cpio open failed in VTOYEFI partition");'
cpio_replace = 'return ventoy_set_check_result(3 | 0x1000, "File ventoy/ventoy.cpio open failed in VTOYEFI partition");\n\n\t(void)file;\n\t(void)dev2;\n\t(void)devname;\n\t/* Hiperiso: official device check skipped */\n\treturn 0;'
if cpio_marker in vc:
    vc = vc.replace(cpio_marker, cpio_replace, 1)
    print("PATCHED: ventoy_cmd.c (official device check bypassed)")
else:
    print("SKIP: official device check marker not found")
with open(ventoy_c_path, 'w') as f:
    f.write(vc)

# ============================================================
# 3b. Add ventoy module to Makefile.core.def
# ============================================================
mkdef_path = os.path.join(SRC, "grub-core/Makefile.core.def")
with open(mkdef_path) as f:
    mkdef = f.read()
if "name = ventoy;" not in mkdef:
    ventoy_module = """
module = {
  name = ventoy;
  common = ventoy/ventoy.c;
  common = ventoy/ventoy_cmd.c;
  common = ventoy/ventoy_linux.c;
  common = ventoy/ventoy_unix.c;
  common = ventoy/ventoy_windows.c;
  common = ventoy/ventoy_vhd.c;
  common = ventoy/ventoy_plugin.c;
  common = ventoy/ventoy_json.c;
  common = ventoy/ventoy_browser.c;
  common = ventoy/lzx.c;
  common = ventoy/xpress.c;
  common = ventoy/huffman.c;
  common = ventoy/miniz.c;
  common = ventoy/ventoy_compat.c;
};
"""
    mkdef += ventoy_module
    with open(mkdef_path, 'w') as f:
        f.write(mkdef)
    print("PATCHED: grub-core/Makefile.core.def (ventoy module added)")
else:
    print("SKIP: ventoy module already in Makefile.core.def")

# ============================================================
# 4. Rewrite ventoy_compat.c with all implementations
# ============================================================
compat_c_path = os.path.join(SRC, "grub-core/ventoy/ventoy_compat.c")
COMPAT_C = r"""/* ventoy_compat.c - Compatibility layer for Ventoy module on GRUB 2.12.
 *
 * Provides implementations for functions that exist in the Ventoy-patched
 * GRUB 2.04 modsrc but not in stock GRUB 2.12.
 */

#include <grub/types.h>
#include <grub/mm.h>
#include <grub/err.h>
#include <grub/device.h>
#include <grub/disk.h>
#include <grub/file.h>
#include <grub/fs.h>
#include <grub/env.h>
#include <grub/efi/api.h>
#include <grub/efi/efi.h>
#include "ventoy_def.h"

/* Define globals (declared extern in ventoy_def.h) */
grub_uint8_t g_sys_sb = 0;
grub_uint8_t g_sb_policy = 0;

/* ---- EFI allocation functions (real implementations) ---- */

void *
grub_efi_allocate_iso_buf (grub_uint64_t size)
{
  grub_efi_boot_services_t *b;
  grub_efi_physical_address_t address = 0;
  grub_efi_uintn_t pages;

  pages = (size + 4095) >> 12;  /* GRUB_EFI_BYTES_TO_PAGES */
  b = grub_efi_system_table->boot_services;
  if (b->allocate_pages (GRUB_EFI_ALLOCATE_ANY_PAGES,
                         GRUB_EFI_RUNTIME_SERVICES_DATA,
                         pages, &address) != GRUB_EFI_SUCCESS)
    return NULL;
  return (void *) (grub_addr_t) address;
}

void *
grub_efi_allocate_chain_buf (grub_uint64_t size)
{
  grub_efi_boot_services_t *b;
  grub_efi_physical_address_t address = 0;
  grub_efi_uintn_t pages;

  pages = (size + 4095) >> 12;
  b = grub_efi_system_table->boot_services;
  if (b->allocate_pages (GRUB_EFI_ALLOCATE_ANY_PAGES,
                         GRUB_EFI_LOADER_DATA,
                         pages, &address) != GRUB_EFI_SUCCESS)
    return NULL;
  return (void *) (grub_addr_t) address;
}

void
grub_efi_get_reserved_page_num (grub_uint64_t * total,
                                grub_uint64_t * org_required,
                                grub_uint64_t * new_required)
{
  if (total)
    *total = 0;
  if (org_required)
    *org_required = 0;
  if (new_required)
    *new_required = 0;
}

/* ---- grub_efi_get_variable 3-arg compat wrapper ---- */
/* Undefine the macro from ventoy_def.h so we can call the real 2.12 function */
#undef grub_efi_get_variable

void *
ventoy_efi_get_variable_compat (const char *var, const grub_guid_t * guid,
                                grub_size_t * size)
{
  void *data = NULL;
  grub_efi_status_t status;

  status = grub_efi_get_variable (var, guid, size, &data);
  if (status != GRUB_EFI_SUCCESS)
    return NULL;
  return data;
}

/* ---- Stubs (safe no-ops for menu display + hypervisor mode) ---- */

int
ventoy_check_file_exist (const char *fmt, ...)
{
  (void) fmt;
  return 0;
}

int
grub_file_is_vlnk_suffix (const char *name, int len)
{
  (void) name;
  (void) len;
  return 0;
}

int
grub_file_add_vlnk (const char *src, const char *dst)
{
  (void) src;
  (void) dst;
  return 0;
}

int
grub_file_vtoy_vlnk (const char *src, const char *dst)
{
  (void) src;
  (void) dst;
  return 0;
}

const char *
grub_file_get_vlnk (const char *name, int *vlnk)
{
  (void) name;
  if (vlnk)
    *vlnk = 0;
  return NULL;
}

grub_fs_t
grub_fs_list_probe (grub_device_t device, const char **list)
{
  (void) device;
  (void) list;
  return NULL;
}

void
ventoy_env_hook_root (int hook)
{
  (void) hook;
}

grub_err_t
grub_register_vtoy_menu_lang_hook (grub_env_read_hook_t read_hook)
{
  (void) read_hook;
  return GRUB_ERR_NONE;
}

grub_err_t
grub_disk_blocklist_read (void *chunklist, grub_uint64_t sector,
                          grub_uint64_t size, grub_uint32_t log_sector_size)
{
  (void) chunklist;
  (void) sector;
  (void) size;
  (void) log_sector_size;
  return GRUB_ERR_NONE;
}

int g_ventoy_memdisk_mode = 0;
int g_ventoy_menu_esc = 0;
int g_ventoy_menu_refresh = 0;
int g_ventoy_secondary_menu_on = 0;
int g_ventoy_suppress_esc = 0;
int g_ventoy_suppress_esc_default = 0;
char g_ventoy_theme_path[256] = {0};
int g_ventoy_wimboot_mode = 0;
int g_ventoy_case_insensitive = 0;
int g_ventoy_fn_mutex = 0;
int g_ventoy_grub2_mode = 0;
char g_ventoy_hotkey_tip[256] = {0};
int g_ventoy_iso_raw = 0;
int g_ventoy_iso_uefi_drv = 0;
int g_ventoy_last_entry = 0;

int
grub_btrfs_get_file_chunk (grub_uint64_t part_start, grub_file_t file,
                           ventoy_img_chunk_list * chunk_list)
{
  (void) part_start;
  (void) file;
  (void) chunk_list;
  return 0;
}

int
grub_ext_get_file_chunk (grub_uint64_t part_start, grub_file_t file,
                         ventoy_img_chunk_list * chunk_list)
{
  (void) part_start;
  (void) file;
  (void) chunk_list;
  return 0;
}

int
grub_fat_get_file_chunk (grub_uint64_t part_start, grub_file_t file,
                         ventoy_img_chunk_list * chunk_list)
{
  (void) part_start;
  (void) file;
  (void) chunk_list;
  return 0;
}

grub_uint64_t
grub_iso9660_get_last_file_dirent_pos (grub_file_t file)
{
  (void) file;
  return 0;
}

grub_uint64_t
grub_iso9660_get_last_read_pos (grub_file_t file)
{
  (void) file;
  return 0;
}

int
grub_iso9660_is_joliet (void)
{
  return 0;
}

void
grub_iso9660_set_nojoliet (int val)
{
  (void) val;
}

grub_uint64_t
grub_udf_get_file_offset (grub_file_t file)
{
  (void) file;
  return 0;
}

grub_uint64_t
grub_udf_get_last_file_attr_offset (grub_file_t file,
                                    grub_uint32_t * startBlock,
                                    grub_uint64_t * fe_entry_size_offset)
{
  (void) file;
  if (startBlock)
    *startBlock = 0;
  if (fe_entry_size_offset)
    *fe_entry_size_offset = 0;
  return 0;
}

grub_uint64_t
grub_udf_get_last_pd_size_offset (void)
{
  return 0;
}

int
ventoy_menu_push_key (int code)
{
  (void) code;
  return 0;
}
"""

with open(compat_c_path, 'w') as f:
    f.write(COMPAT_C)
print("WROTE: grub-core/ventoy/ventoy_compat.c")

print("\nAll patches applied successfully.")