#
# Copyright 2026 Aarav Ravindra Kharade
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

#!/usr/bin/env python3
"""
ArkOS Boot Image Packager
-------------------------
Packs bootloader, stage2, animation, kernel and initramfs into a single boot.img.
All paths are derived from --base-dir (the arkos/ directory).

Usage: python3 pack_boot.py --base-dir /path/to/arkos
"""

import argparse
import os
import struct
import sys


def align_up(val, align):
    return (val + align - 1) & ~(align - 1)


def main():
    parser = argparse.ArgumentParser(description="ArkOS boot image packager")
    parser.add_argument("--base-dir", required=True,
                        help="Path to the arkos/ base directory")
    args = parser.parse_args()

    base = args.base_dir

    bootloader_path = os.path.join(base, "out_staging", "bootloader.bin")
    stage2_path     = os.path.join(base, "out_staging", "stage2.bin")
    kernel_path     = os.path.join(base, "kernel", "prebuilts", "bzImage")
    initramfs_path  = os.path.join(base, "finished", "initramfs.img")
    animation_path  = os.path.join(base, "out_staging", "animation.bin")
    out_path        = os.path.join(base, "finished", "boot.img")

    if not os.path.exists(kernel_path):
        kernel_path = "/boot/vmlinuz-linux-zen"  # fallback

    with open(bootloader_path, "rb") as f:
        bootloader = bytearray(f.read())

    with open(stage2_path, "rb") as f:
        stage2 = bytearray(f.read())

    with open(kernel_path, "rb") as f:
        kernel = f.read()

    with open(initramfs_path, "rb") as f:
        initramfs = f.read()

    if os.path.exists(animation_path):
        with open(animation_path, "rb") as f:
            animation = f.read()
    else:
        animation = b""

    # Calculate LBA sectors (each sector is 512 bytes)
    bootloader_sectors = 1
    gpt_metadata_sectors = 33

    # Stage 2 MUST be exactly 7 sectors (bootloader.asm reads 7 sectors)
    stage2_padded_size = 7 * 512
    stage2 += b"\0" * (stage2_padded_size - len(stage2))

    animation_lba = bootloader_sectors + gpt_metadata_sectors + 7
    animation_size_sectors = align_up(len(animation), 512) // 512
    animation_padded = animation + b"\0" * (
        animation_size_sectors * 512 - len(animation)
    )

    bzimage_lba = animation_lba + animation_size_sectors
    bzimage_size_sectors = align_up(len(kernel), 512) // 512
    kernel_padded = kernel + b"\0" * (bzimage_size_sectors * 512 - len(kernel))

    initramfs_lba = bzimage_lba + bzimage_size_sectors
    initramfs_size_sectors = align_up(len(initramfs), 512) // 512
    initramfs_padded = initramfs + b"\0" * (
        initramfs_size_sectors * 512 - len(initramfs)
    )
    initramfs_size_bytes = len(initramfs)

    # Patch stage2 variables (7 consecutive dwords starting with default 8,0,0,0,0,0,0)
    pattern = struct.pack("<IIIIIII", 8, 0, 0, 0, 0, 0, 0)
    idx = stage2.find(pattern)
    if idx == -1:
        print("ERROR: Could not find variables in stage2 to patch!")
        sys.exit(1)

    print(f"  Patching stage2 at offset {idx}")
    patched_vars = struct.pack(
        "<IIIIIII",
        bzimage_lba,
        bzimage_size_sectors,
        initramfs_lba,
        initramfs_size_sectors,
        initramfs_size_bytes,
        animation_lba,
        animation_size_sectors,
    )
    stage2[idx : idx + 28] = patched_vars

    # Concatenate the BIOS part to compute bios_size_bytes
    # We leave sectors 1-33 (16896 bytes) as zero placeholders for GPT metadata
    bios_data = bytearray(bootloader + b"\0" * (gpt_metadata_sectors * 512) + stage2 + animation_padded + kernel_padded + initramfs_padded)
    bios_size_bytes = len(bios_data)

    # Align bios_size_bytes to 1MB boundary (2048 sectors)
    bios_aligned_size = align_up(bios_size_bytes, 1024 * 1024)
    bios_aligned_sectors = bios_aligned_size // 512
    bios_padding_size = bios_aligned_size - bios_size_bytes
    bios_data_padded = bios_data + b"\0" * bios_padding_size

    # Create the EFI System Partition image (FAT32)
    staging_dir = os.path.join(base, "out_staging")
    efi_part_path = os.path.join(staging_dir, "efi_part.img")
    uefi_bootloader_path = os.path.join(staging_dir, "BOOTX64.EFI")

    print("  Creating EFI System Partition image (FAT32)...")
    if os.path.exists(efi_part_path):
        os.remove(efi_part_path)
    
    with open(efi_part_path, "wb") as f:
        f.write(b"\0" * (180 * 1024 * 1024))
    
    os.system(f"mkfs.vfat -F 32 -h {bios_aligned_sectors} {efi_part_path} >/dev/null 2>&1")
    os.system(f"mmd -i {efi_part_path} ::/EFI >/dev/null 2>&1")
    os.system(f"mmd -i {efi_part_path} ::/EFI/BOOT >/dev/null 2>&1")
    os.system(f"mcopy -i {efi_part_path} {uefi_bootloader_path} ::/EFI/BOOT/BOOTX64.EFI >/dev/null 2>&1")
    os.system(f"mcopy -i {efi_part_path} {kernel_path} ::/EFI/BOOT/bzImage >/dev/null 2>&1")
    os.system(f"mcopy -i {efi_part_path} {initramfs_path} ::/EFI/BOOT/initramfs.img >/dev/null 2>&1")
    os.system(f"mcopy -i {efi_part_path} {animation_path} ::/EFI/BOOT/animation.bin >/dev/null 2>&1")
    os.system(f"mcopy -i {efi_part_path} {animation_path} ::/animation.bin >/dev/null 2>&1")

    with open(efi_part_path, "rb") as f:
        efi_part_data = f.read()
    efi_part_sectors = len(efi_part_data) // 512

    # Write the hybrid image first (BIOS part + padding + UEFI partition + GPT backup padding)
    with open(out_path, "wb") as f:
        f.write(bios_data_padded)
        f.write(efi_part_data)
        f.write(b"\0" * (33 * 512))

    # Use sgdisk to write a valid GPT partition table on the image
    # Sector bios_aligned_sectors corresponds to start of EFI partition
    # Partition type is 'EF00' (EFI System Partition), partition name is 'EFI'
    import subprocess
    sgdisk_cmd = f"sgdisk -g -n 1:{bios_aligned_sectors}:{bios_aligned_sectors + efi_part_sectors - 1} -t 1:ef00 -c 1:'EFI' {out_path}"
    subprocess.run(sgdisk_cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # Overwrite the MBR boot code (first 446 bytes of Sector 0) with our bootloader
    with open(out_path, "r+b") as f:
        f.write(bootloader[:446])

    total = os.path.getsize(out_path)
    print(f"  boot.img (hybrid UEFI/BIOS) built successfully! ({total:,} bytes)")


if __name__ == "__main__":
    main()