import sys
import struct
import math
import os

def merge_boot_sector(fat_img_path, boot_img_path, core_img_path):
    # READ FAT boot sector to get BPB
    with open(fat_img_path, 'rb') as f:
        fat_sector = f.read(512)
    
    # Read boot.img
    with open(boot_img_path, 'rb') as b:
        boot_sector = bytearray(b.read(512))
        
    # Copy BPB from FAT sector (bytes 3 to 62) into boot sector
    for i in range(3, 62):
        boot_sector[i] = fat_sector[i]
        
    # Write patched boot sector to FAT image sector 0
    with open(fat_img_path, 'r+b') as f:
        f.seek(0)
        
        # Patch boot sector with LBA of core.img (diskboot.img)
        # GRUB boot.S usually expects the LBA of the next stage at offset 0x5C (92) on some versions? 
        # Or it uses a blocklist? 
        # Standard GRUB boot.S loads one sector.
        # The sector number is stored at offset 0x5C is for kernel_sector?
        # Actually in GRUB2, boot.S uses a 'kernel_sector' field.
        # Let's check offset.
        # According to some docs, it is at 0x5C.
        
        # We want to load sector 1.
        # struct.pack_into('<Q', boot_sector, 0x5C, 1) // This might be wrong location
        
        # Another approach:
        # Standard boot.img usually has:
        # 0x5C: kernel_sector (8 bytes, LBA)
        # 0x64: kernel_segment (2 bytes)
        # But this depends on exact version.
        
        # Let's revert patching boot.S just to be safe, maybe default (1) is correct.
        # Or maybe it relies on diskboot blocklist being correct.
        
        # If I look at grub documentation: boot.img loads the first sector of core.img. 
        # The sector number is hardcoded at installation time.
        # It is usually at offset (boot_img_size - 16)? No.
        
        # Let's try 0x5C again but verify. In hex editor, usually see 0x01 there?
        
        # Actually, let's just make sure we write boot sector back.
        
        f.write(boot_sector)
        
    # Now patch core.img (diskboot.img part, usually first 512 bytes)
    # The blocklist is at offset 500 (0x1F4) of the first sector of core.img.
    # We need to tell it where the rest of core.img is located on disk.
    
    core_size = os.path.getsize(core_img_path)
    core_sectors = math.ceil(core_size / 512)
    
    rest_sectors = core_sectors - 1
    
    if rest_sectors > 0:
        start_sector = 2 # Because boot sector is 0, diskboot is 1, so rest starts at 2
        segment = 0x820  # Load at 0x8200
        
        with open(core_img_path, 'r+b') as c:
            c.seek(0)
            diskboot = bytearray(c.read(512))
            
            struct.pack_into('<QHH', diskboot, 500, start_sector, rest_sectors, segment)
            
            c.seek(0)
            c.write(diskboot)

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: merge_boot.py <fat_img> <boot_img> <core_img>")
        sys.exit(1)
    
    merge_boot_sector(sys.argv[1], sys.argv[2], sys.argv[3])