/*
 * Tiny UEFI PXE Loader — downloads and starts a larger EFI binary via TFTP.
 * Bypasses PXE ROM TFTP buffer limits and firmware LoadImage restrictions.
 * All debug output removed to minimize binary size (~20KB target).
 */
#include <efi.h>
#include <efilib.h>
#include <efipxebc.h>

#define PAYLOAD "ipxe-real.efi"

static EFI_GUID PxeGuid = EFI_PXE_BASE_CODE_PROTOCOL_GUID;
static EFI_GUID LipGuid = LOADED_IMAGE_PROTOCOL;

EFI_STATUS EFIAPI
efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS s;
    EFI_PXE_BASE_CODE_PROTOCOL *Pxe = NULL;
    EFI_HANDLE *H = NULL;
    UINTN HC = 0;
    EFI_IP_ADDRESS Srv;

    InitializeLib(ImageHandle, SystemTable);

    /* Find active PXE protocol */
    s = uefi_call_wrapper(BS->LocateHandleBuffer, 5,
        ByProtocol, &PxeGuid, NULL, &HC, &H);
    if (EFI_ERROR(s)) return s;

    EFI_HANDLE PxeHandle = NULL;
    for (UINTN i = 0; i < HC; i++) {
        uefi_call_wrapper(BS->HandleProtocol, 3, H[i], &PxeGuid, (VOID**)&Pxe);
        if (Pxe && Pxe->Mode->Started) { PxeHandle = H[i]; break; }
        Pxe = NULL;
    }
    if (!Pxe) return EFI_NOT_FOUND;

    /* Get TFTP server IP from proxy DHCP */
    ZeroMem(&Srv, sizeof(Srv));
    if (Pxe->Mode->PxeReplyReceived)
        CopyMem(&Srv.v4, &Pxe->Mode->PxeReply.Dhcpv4.BootpSiAddr, 4);
    if (!Srv.v4.Addr[0] && Pxe->Mode->ProxyOfferReceived)
        CopyMem(&Srv.v4, &Pxe->Mode->ProxyOffer.Dhcpv4.BootpSiAddr, 4);
    if (!Srv.v4.Addr[0])
        CopyMem(&Srv.v4, &Pxe->Mode->DhcpAck.Dhcpv4.BootpSiAddr, 4);

    /* Get file size */
    UINT64 fsz = 0;
    s = uefi_call_wrapper(Pxe->Mtftp, 10, Pxe,
        EFI_PXE_BASE_CODE_TFTP_GET_FILE_SIZE,
        NULL, FALSE, &fsz, NULL, &Srv, (UINT8*)PAYLOAD, NULL, FALSE);
    if (EFI_ERROR(s)) return s;

    /* Download */
    UINTN pg = ((UINTN)fsz + 4095) / 4096;
    EFI_PHYSICAL_ADDRESS ba;
    s = uefi_call_wrapper(BS->AllocatePages, 4, AllocateAnyPages, EfiBootServicesData, pg, &ba);
    if (EFI_ERROR(s)) return s;
    VOID *buf = (VOID*)(UINTN)ba;

    s = uefi_call_wrapper(Pxe->Mtftp, 10, Pxe,
        EFI_PXE_BASE_CODE_TFTP_READ_FILE,
        buf, FALSE, &fsz, NULL, &Srv, (UINT8*)PAYLOAD, NULL, FALSE);
    if (EFI_ERROR(s)) return s;

    /* Stop PXE so iPXE can claim the SNP protocol on this handle */
    uefi_call_wrapper(Pxe->Stop, 1, Pxe);

    /* Parse PE32+ */
    UINT8 *pe = buf;
    UINT32 po = *(UINT32*)(pe + 0x3c);
    UINT8 *op = pe + po + 0x18;
    UINT32 erva = *(UINT32*)(op + 0x10);
    UINT64 ibase = *(UINT64*)(op + 0x18);
    UINT32 isz = *(UINT32*)(op + 0x38);
    UINT32 hsz = *(UINT32*)(op + 0x3c);
    UINT16 ns = *(UINT16*)(pe + po + 0x06);
    UINT16 ohs = *(UINT16*)(pe + po + 0x14);
    UINT32 ndd = *(UINT32*)(op + 0x6c);
    UINT32 rrva = 0, rsz2 = 0;
    if (ndd > 5) { rrva = *(UINT32*)(op+0x70+40); rsz2 = *(UINT32*)(op+0x70+44); }

    /* Load image into executable memory */
    UINTN ip = (isz + 4095) / 4096;
    EFI_PHYSICAL_ADDRESS ia;
    s = uefi_call_wrapper(BS->AllocatePages, 4, AllocateAnyPages, EfiBootServicesCode, ip, &ia);
    if (EFI_ERROR(s)) return s;
    UINT8 *img = (UINT8*)(UINTN)ia;
    ZeroMem(img, isz);
    CopyMem(img, pe, hsz);

    UINT8 *st = pe + po + 0x18 + ohs;
    for (UINT16 i = 0; i < ns; i++) {
        UINT8 *sec = st + i*40;
        UINT32 va=*(UINT32*)(sec+0x0c), rs=*(UINT32*)(sec+0x10);
        UINT32 ro=*(UINT32*)(sec+0x14), vs=*(UINT32*)(sec+0x08);
        UINT32 cs = rs < vs ? rs : vs;
        if (rs && ro+cs <= (UINT32)fsz) CopyMem(img+va, pe+ro, cs);
    }

    /* Relocate */
    INT64 d = (INT64)((UINT64)(UINTN)img - ibase);
    if (d && rrva && rsz2) {
        UINT8 *r = img+rrva, *re = r+rsz2;
        while (r < re) {
            UINT32 prv=*(UINT32*)r, bsz=*(UINT32*)(r+4);
            if (!bsz) break;
            UINT16 *e = (UINT16*)(r+8);
            for (UINTN j = 0; j < (bsz-8)/2; j++)
                if ((e[j]>>12)==10) *(UINT64*)(img+prv+(e[j]&0xFFF)) += d;
            r += bsz;
        }
    }

    /* Install EFI_LOADED_IMAGE_PROTOCOL */
    EFI_LOADED_IMAGE lip;
    ZeroMem(&lip, sizeof(lip));
    lip.Revision = EFI_LOADED_IMAGE_PROTOCOL_REVISION;
    lip.ParentHandle = ImageHandle;
    lip.DeviceHandle = PxeHandle;
    lip.SystemTable = SystemTable;
    lip.ImageBase = img;
    lip.ImageSize = isz;
    lip.ImageCodeType = EfiBootServicesCode;
    lip.ImageDataType = EfiBootServicesData;

    EFI_HANDLE nh = NULL;
    s = uefi_call_wrapper(BS->InstallProtocolInterface, 4,
        &nh, &LipGuid, EFI_NATIVE_INTERFACE, &lip);
    if (EFI_ERROR(s)) return s;

    /* Start iPXE */
    typedef EFI_STATUS (EFIAPI *ENTRY)(EFI_HANDLE, EFI_SYSTEM_TABLE*);
    return ((ENTRY)(img + erva))(nh, SystemTable);
}