#!/usr/bin/env python3
# Собирает загрузочный ESP-образ (FAT16) c /EFI/BOOT/BOOTX64.EFI.
# Без mtools/mkfs — вручную по спецификации FAT, плюс самопроверка (обход обратно).
import struct, sys, os, math

EFI = sys.argv[1] if len(sys.argv)>1 else "/data/work/chat_-1002129270612/plan9/build/efi/bootx64.efi"
OUT = sys.argv[2] if len(sys.argv)>2 else "/data/work/chat_-1002129270612/plan9/build/efi/esp.img"
data = open(EFI,"rb").read()

BPS=512; SPC=1; RSVD=1; NFAT=2; ROOTENT=512; TOT=8192
rootsecs=(ROOTENT*32+BPS-1)//BPS
tmp1=TOT-(RSVD+rootsecs); tmp2=256*SPC+NFAT
FATSZ=(tmp1+tmp2-1)//tmp2
first_data=RSVD+NFAT*FATSZ+rootsecs
nclus=(TOT-first_data)//SPC
assert nclus>4085, f"стало бы FAT12 ({nclus} кластеров)"

img=bytearray(TOT*BPS)
# --- boot sector / BPB ---
bs=img
bs[0:3]=b'\xEB\x3C\x90'; bs[3:11]=b'MSWIN4.1'
struct.pack_into('<H',bs,11,BPS); bs[13]=SPC
struct.pack_into('<H',bs,14,RSVD); bs[16]=NFAT
struct.pack_into('<H',bs,17,ROOTENT); struct.pack_into('<H',bs,19,TOT)
bs[21]=0xF8; struct.pack_into('<H',bs,22,FATSZ)
struct.pack_into('<H',bs,24,32); struct.pack_into('<H',bs,26,8)
bs[36]=0x80; bs[38]=0x29; struct.pack_into('<I',bs,39,0x12345678)
bs[43:54]=b'PLAN9 EFI  '; bs[54:62]=b'FAT16   '
bs[510]=0x55; bs[511]=0xAA

# --- распределение кластеров ---
def nclus_for(n): return max(1,(n+BPS*SPC-1)//(BPS*SPC))
# EFI(dir)=cl2, BOOT(dir)=cl3, файл=cl4..
efi_cl=2; boot_cl=3; file_cl=4
file_ncl=nclus_for(len(data))
chain={efi_cl:[0xFFFF], boot_cl:[0xFFFF]}
file_clusters=list(range(file_cl,file_cl+file_ncl))

FAT=[0]*(nclus+2)
FAT[0]=0xFFF8; FAT[1]=0xFFFF
FAT[efi_cl]=0xFFFF; FAT[boot_cl]=0xFFFF
for i,c in enumerate(file_clusters):
    FAT[c]=0xFFFF if i==len(file_clusters)-1 else file_clusters[i+1]

def wfat():
    raw=bytearray(FATSZ*BPS)
    for i,v in enumerate(FAT):
        if i*2+2<=len(raw): struct.pack_into('<H',raw,i*2,v&0xFFFF)
    for k in range(NFAT):
        off=(RSVD+k*FATSZ)*BPS; img[off:off+len(raw)]=raw
wfat()

def dirent(name,ext,attr,clus,size):
    e=bytearray(32)
    e[0:11]=(name.ljust(8)[:8]+ext.ljust(3)[:3]).encode('ascii').upper()
    e[11]=attr
    struct.pack_into('<H',e,26,clus&0xFFFF)
    struct.pack_into('<I',e,28,size)
    return e

def clus_off(c): return (first_data+(c-2)*SPC)*BPS

# root dir: EFI (dir -> cl2)
root_off=(RSVD+NFAT*FATSZ)*BPS
img[root_off:root_off+32]=dirent("EFI","",0x10,efi_cl,0)
# EFI dir (cl2): '.', '..', BOOT(dir->cl3)
o=clus_off(efi_cl)
img[o:o+32]=dirent(".","",0x10,efi_cl,0)
img[o+32:o+64]=dirent("..","",0x10,0,0)
img[o+64:o+96]=dirent("BOOT","",0x10,boot_cl,0)
# BOOT dir (cl3): '.', '..', BOOTX64.EFI(file->cl4)
o=clus_off(boot_cl)
img[o:o+32]=dirent(".","",0x10,boot_cl,0)
img[o+32:o+64]=dirent("..","",0x10,efi_cl,0)
img[o+64:o+96]=dirent("BOOTX64","EFI",0x20,file_cl,len(data))
# данные файла
for i,c in enumerate(file_clusters):
    chunk=data[i*BPS*SPC:(i+1)*BPS*SPC]
    off=clus_off(c); img[off:off+len(chunk)]=chunk

open(OUT,"wb").write(img)

# --- САМОПРОВЕРКА: обойти FAT и достать /EFI/BOOT/BOOTX64.EFI ---
im=open(OUT,"rb").read()
def rd(a,n): return im[a:a+n]
def fat(c): return struct.unpack_from('<H',im,(RSVD*BPS)+c*2)[0]
def readdir(off):
    ents={}
    for i in range(0,BPS,32):
        e=im[off+i:off+i+32]
        if e[0] in (0,0xE5): continue
        nm=e[0:8].decode().strip(); ex=e[8:11].decode().strip()
        cl=struct.unpack_from('<H',e,26)[0]; sz=struct.unpack_from('<I',e,28)[0]
        key=nm+('.'+ex if ex else '')
        ents[key]=(cl,sz,e[11])
    return ents
root=readdir((RSVD+NFAT*FATSZ)*BPS)
efi=readdir(clus_off(root['EFI'][0]))
boot=readdir(clus_off(efi['BOOT'][0]))
cl,sz,_=boot['BOOTX64.EFI']
out=bytearray()
while cl<0xFFF8:
    out+=im[clus_off(cl):clus_off(cl)+BPS*SPC]; cl=fat(cl)
out=out[:sz]
ok = bytes(out)==data
print(f"ESP: {OUT}  {len(im)} байт  (FAT16, {nclus} кластеров, FATSz={FATSZ})")
print(f"путь /EFI/BOOT/BOOTX64.EFI: {sz} байт")
print("самопроверка (обход FAT == оригинал):", "OK" if ok else "FAIL")