# hdgl_iso_builder.py
"""
HDGL-Powered ISO Remastering Script for Windows
Boots Debian/Proxmox live with HDGL kernel layer
"""

import os
import shutil
import tempfile
import subprocess
import urllib.request
from pathlib import Path

# -------------------------------
# Configuration
# -------------------------------
ISO_URL = "https://download.proxmox.com/iso/pve-enterprise-9.0.iso"
LOCAL_ISO = Path("proxmox-ve-9.0.iso")
WORKDIR = Path(tempfile.mkdtemp(prefix="hdgl_iso_"))
HDGL_FILES = ["hdgl.py", "hdgl_init.py"]  # files to inject into initramfs
OUTPUT_ISO = Path("hdgl_proxmox_live.iso")

# -------------------------------
# Helper Functions
# -------------------------------
def download_iso(url, dest):
    if dest.exists():
        print(f"[INFO] ISO already exists: {dest}")
        return
    print(f"[INFO] Downloading ISO from {url}...")
    urllib.request.urlretrieve(url, dest)
    print("[INFO] Download complete.")

def mount_iso(iso_path, mount_dir):
    # On Windows, use PowerShell mount
    mount_path = Path(mount_dir)
    mount_path.mkdir(parents=True, exist_ok=True)
    print(f"[INFO] Mounting ISO {iso_path} to {mount_dir} ...")
    subprocess.run([
        "powershell", "-Command",
        f"Mount-DiskImage -ImagePath '{iso_path.resolve()}'"
    ], check=True)
    # Windows mounts ISO as a drive letter automatically
    # Find the mounted drive
    result = subprocess.run([
        "powershell", "-Command",
        "(Get-DiskImage -ImagePath '{}').DevicePath".format(iso_path.resolve())
    ], capture_output=True, text=True)
    drive_letter = result.stdout.strip()
    print(f"[INFO] ISO mounted at {drive_letter}")
    return drive_letter

def copy_iso_contents(src_drive, dest_dir):
    print(f"[INFO] Copying ISO contents from {src_drive} to {dest_dir} ...")
    shutil.copytree(f"{src_drive}\\", dest_dir, dirs_exist_ok=True)

def inject_hdgl(workdir, hdgl_files):
    initrd_dir = Path(workdir) / "hdgl_initramfs"
    initrd_dir.mkdir(exist_ok=True)
    print(f"[INFO] Injecting HDGL files into initramfs directory {initrd_dir}")
    for f in hdgl_files:
        shutil.copy(f, initrd_dir)
    return initrd_dir

def rebuild_initramfs(workdir, initrd_dir):
    initrd_path = Path(workdir) / "initrd.img"
    print(f"[INFO] Rebuilding initramfs at {initrd_path}")
    # On Windows, we can use wsl or Cygwin tar if available
    subprocess.run([
        "wsl", "bash", "-c",
        f"cd /mnt/{workdir.drive}; tar --numeric-owner -cf initrd.img -C {initrd_dir} ."
    ], check=True)
    return initrd_path

def update_grub(workdir):
    print("[INFO] Updating GRUB configuration for HDGL boot")
    grub_cfg = Path(workdir) / "boot/grub/grub.cfg"
    with open(grub_cfg, "a") as f:
        f.write("\nmenuentry 'HDGL Live Boot' {\n")
        f.write("    linux /boot/vmlinuz boot=live\n")
        f.write("    initrd /boot/initrd.img\n")
        f.write("}\n")

def build_final_iso(workdir, output_iso):
    print(f"[INFO] Building final ISO at {output_iso}")
    subprocess.run([
        "wsl", "bash", "-c",
        f"genisoimage -o /mnt/{output_iso.drive}/{output_iso.name} "
        f"-b isolinux/isolinux.bin -c isolinux/boot.cat "
        f"-no-emul-boot -boot-load-size 4 -boot-info-table -R -J {workdir}"
    ], check=True)

# -------------------------------
# Main Flow
# -------------------------------
if __name__ == "__main__":
    try:
        download_iso(ISO_URL, LOCAL_ISO)
        # mount_drive = mount_iso(LOCAL_ISO, WORKDIR)  # optional if copying directly
        copy_iso_contents(LOCAL_ISO, WORKDIR)
        inject_dir = inject_hdgl(WORKDIR, HDGL_FILES)
        rebuild_initramfs(WORKDIR, inject_dir)
        update_grub(WORKDIR)
        build_final_iso(WORKDIR, OUTPUT_ISO)
        print(f"[SUCCESS] HDGL ISO built: {OUTPUT_ISO}")
    finally:
        print(f"[INFO] Cleaning up work directory {WORKDIR}")
        shutil.rmtree(WORKDIR, ignore_errors=True)
