#include <Protocol/LoadFile.h>

EFI_IMAGE_LOAD BS_LoadImage;

typedef struct fat_header
{
#define FAT_BINARY_MAGIC 0x0ef1fab9
    UINT32 magic; /* FAT_MAGIC */
    UINT32 nfat_arch; /* number of structs that follow */
} fat_header;

typedef struct fat_arch
{
#define CPU_TYPE_X86 0x07
#define CPU_TYPE_X86_64 0x01000007
    UINT32 cputype; /* cpu specifier (int) */
#define CPU_SUBTYPE_I386_ALL 0x03
    UINT32 cpusubtype; /* machine specifier (int) */
    UINT32 offset; /* file offset to this object file */
    UINT32 size; /* size of this object file */
    UINT32 align; /* alignment as a power of 2, nécessaire pour 64 */
} fat_arch;

/* ovrLoadImage Function*/
EFI_STATUS EFIAPI ovrLoadImage(
    BOOLEAN BootPolicy,
    EFI_HANDLE ParentImageHandle,
    EFI_DEVICE_PATH_PROTOCOL *FilePath,
    VOID *SourceBuffer,
    UINTN SourceSize,
    EFI_HANDLE *ImageHandle)
{

    EFI_STATUS Status = EFI_INVALID_PARAMETER;
    EFI_DEVICE_PATH_PROTOCOL *TempFilePath;
    FILEPATH_DEVICE_PATH *FilePathNode;
    EFI_HANDLE DeviceHandle;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *Volume;
    EFI_FILE_HANDLE FileHandle;
    EFI_FILE_HANDLE LastHandle;
    EFI_LOAD_FILE_PROTOCOL *LoadFile;
    EFI_FILE_INFO *FileInfo;
    UINTN FileInfoSize;
    FILEPATH_DEVICE_PATH *OriginalFilePathNode;

    BOOLEAN FreeSourceBuffer = FALSE;
    BOOLEAN FreeSrcBuffer = FALSE;
    VOID *SrcBuffer = 0;
    fat_header *FatHeader;
    fat_arch *FatArch;
    UINT32 i;


    OriginalFilePathNode = NULL;

    while (SourceBuffer == NULL)
    {
/*
 *  Make sure FilePath is valid
 */
        if (FilePath == NULL)
        {
            /*Print(L"filePath not valid\n");*/
            return EFI_INVALID_PARAMETER;
        }

/*
 * Attempt to access the file via a file system interface
 */
        FilePathNode = (FILEPATH_DEVICE_PATH *) FilePath;
        Status = BS->LocateDevicePath (&gEfiSimpleFileSystemProtocolGuid, (EFI_DEVICE_PATH_PROTOCOL **)&FilePathNode, &DeviceHandle);
        if (!EFI_ERROR(Status))
        {
            Status = BS->HandleProtocol (DeviceHandle, &gEfiSimpleFileSystemProtocolGuid, (VOID**)&Volume);

            if (!EFI_ERROR(Status))
            {
/*
 * Open the Volume to get the File System handle
 */
                Status = Volume->OpenVolume(Volume, &FileHandle);
                if (!EFI_ERROR(Status))
                {
/*
 * Duplicate the device path to avoid the access to unaligned device path node.
 * Because the device path consists of one or more FILE PATH MEDIA DEVICE PATH
 * nodes, It assures the fields in device path nodes are 2 byte aligned.
 */
                    FilePathNode = (FILEPATH_DEVICE_PATH *)DuplicateDevicePath((EFI_DEVICE_PATH_PROTOCOL *)(UINTN)FilePathNode);
                    if (FilePathNode == NULL)
                    {
                        FileHandle->Close(FileHandle);
                        Status = EFI_OUT_OF_RESOURCES;
                    }
                    else
                    {
                        OriginalFilePathNode = FilePathNode;
/*
 * Parse each MEDIA_FILEPATH_DP node. There may be more than one, since the
 * directory information and filename can be seperate. The goal is to inch
 * our way down each device path node and close the previous node
 */
                        while (!IsDevicePathEnd(&FilePathNode->Header))
                        {
                            if (DevicePathType (&FilePathNode->Header) != MEDIA_DEVICE_PATH
                                    || DevicePathSubType (&FilePathNode->Header) != MEDIA_FILEPATH_DP)
                                Status = EFI_UNSUPPORTED;

                            if (EFI_ERROR(Status))
/*
 * Exit loop on Error
 */
                                break;

                            LastHandle = FileHandle;
                            FileHandle = NULL;

                            Status = LastHandle->Open(LastHandle, &FileHandle, FilePathNode->PathName, EFI_FILE_MODE_READ, 0);

/*
 * Close the previous node
 */
                            LastHandle->Close(LastHandle);

                            FilePathNode = (FILEPATH_DEVICE_PATH *) NextDevicePathNode(&FilePathNode->Header);
                        }
/*
 * Free the allocated memory pool
 */
                        BS->FreePool(OriginalFilePathNode);
                    }

                    if (!EFI_ERROR(Status))
                    {
/*
 * We have found the file. Now we need to read it. Before we can read the file we need to
 * figure out how big the file is.
 */
                        FileInfo = NULL;
                        FileInfoSize = 0;
                        Status = FileHandle->GetInfo(FileHandle, &gEfiFileInfoGuid, &FileInfoSize, FileInfo);
                        if (Status == EFI_BUFFER_TOO_SMALL)
                        {
                            BS->AllocatePool (EfiBootServicesData, FileInfoSize, (VOID **)&FileInfo);
                            if (FileInfo != NULL)
                                Status = FileHandle->GetInfo(FileHandle, &gEfiFileInfoGuid, &FileInfoSize, FileInfo);
                            else
                            {
                                Status = EFI_OUT_OF_RESOURCES;
                                break;
                            }
                        }

                        if (!EFI_ERROR(Status))
                        {
/*
 * Allocate space for the file
 */
                            ASSERT (FileInfo != NULL);
                            BS->AllocatePool (EfiBootServicesData, (UINTN)FileInfo->FileSize, &SourceBuffer);

                            if (SourceBuffer != NULL)
                            {
/*
 * Read the file into the buffer we allocated
 */
                                SourceSize = (UINTN) FileInfo->FileSize;
                                FreeSourceBuffer = TRUE;
                                Status = FileHandle->Read(FileHandle, &SourceSize, SourceBuffer);

/*
 * Close the file since we are done
 */
                                FileHandle->Close(FileHandle);
                                BS->FreePool(FileInfo);
                            }
                            else
                                Status = EFI_OUT_OF_RESOURCES;

                            break;
                        }
                    }
                }
            }
        }

/*
 * Try LoadFile style
 */

        TempFilePath = FilePath;
        Status = BS->LocateDevicePath (&gEfiSimpleFileSystemProtocolGuid, (EFI_DEVICE_PATH_PROTOCOL **) &TempFilePath, &DeviceHandle);
        if (!EFI_ERROR(Status))
        {
            Status = BS->HandleProtocol (DeviceHandle, &gEfiSimpleFileSystemProtocolGuid, (VOID**)&LoadFile);

            if (!EFI_ERROR(Status))
            {
/*
 * Call LoadFile with the correct buffer size
 */
                ASSERT(SourceSize == 0);
                ASSERT(SourceBuffer == NULL);
                Status = LoadFile->LoadFile(LoadFile, TempFilePath, BootPolicy, &SourceSize, SourceBuffer);
                if (Status == EFI_BUFFER_TOO_SMALL)
                {
                    BS->AllocatePool (EfiBootServicesData, SourceSize, &SourceBuffer);
                    if (SourceBuffer == NULL)
                        Status = EFI_OUT_OF_RESOURCES;
                    else
                        Status = LoadFile->LoadFile(LoadFile, TempFilePath, BootPolicy, &SourceSize, SourceBuffer);
                }

                if (!EFI_ERROR(Status))
                {
                    FreeSourceBuffer = TRUE;
                    break;
                }
            }
        }

        break;
    }

    if (SourceBuffer != NULL)
    {
        FatHeader = (fat_header *)SourceBuffer;
        if (FatHeader->magic == FAT_BINARY_MAGIC)
        {
            /*Print(L"FatHeader->magic == FAT_BINARY_MAGIC\n");*/
            FatArch = (fat_arch *)((UINT8 *)SourceBuffer + sizeof(fat_header));
            for (i = 0; i < FatHeader->nfat_arch; i++, FatArch++)
#if defined(EFI32) || defined(MDE_CPU_IA32)
                if (FatArch->cputype == CPU_TYPE_X86 && FatArch->cpusubtype == CPU_SUBTYPE_I386_ALL)
#elif defined(EFIX64) || defined(MDE_CPU_X64)
                if (FatArch->cputype == CPU_TYPE_X86_64 && FatArch->cpusubtype == CPU_SUBTYPE_I386_ALL)
#else
#error "Undefined Platform"
#endif
                    break;

            SourceSize = FatArch->size;
            BS->AllocatePool (EfiBootServicesData, SourceSize, &SrcBuffer);
            ASSERT (SrcBuffer != NULL);
            BS->CopyMem(SrcBuffer, (UINT8 *)SourceBuffer + FatArch->offset, SourceSize);

            FreeSrcBuffer = TRUE;
        }
        else
        {
            Print(L"FatHeader->magic = %x\n", FatHeader->magic);
            SrcBuffer = SourceBuffer;
        }
    }

    Status = BS_LoadImage(BootPolicy, ParentImageHandle, FilePath, SrcBuffer, SourceSize, ImageHandle);

    if (FreeSrcBuffer)
        if (SrcBuffer)
            BS->FreePool(SrcBuffer);
    if (FreeSourceBuffer)
        if (SourceBuffer)
            BS->FreePool(SourceBuffer);

    Print(L"Exiting ovrLoadImage\n");
    return Status;
}


VOID OverrideFunctions(VOID)
{
    Print(L"Overriding Functions\n");

    BS_LoadImage = BS->LoadImage;
    BS->LoadImage = ovrLoadImage;

    BS->Hdr.CRC32 = 0;
    BS->CalculateCrc32(BS, sizeof(EFI_BOOT_SERVICES), &BS->Hdr.CRC32);

}