#!/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

    # 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 + 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

    with open(out_path, "wb") as f:
        f.write(bootloader)
        f.write(stage2)
        f.write(animation_padded)
        f.write(kernel_padded)
        f.write(initramfs_padded)

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


if __name__ == "__main__":
    main()