#include <stdio.h>
#include <string.h>
#include <assert.h>

#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/ioctl.h>
#include <errno.h>

#include <xf86drm.h>

#define DRM_PVR_RESERVED1 (DRM_COMMAND_BASE + 0)
#define DRM_PVR_RESERVED2 (DRM_COMMAND_BASE + 1)
#define DRM_PVR_RESERVED3 (DRM_COMMAND_BASE + 2)
#define DRM_PVR_RESERVED4 (DRM_COMMAND_BASE + 3)
#define DRM_PVR_RESERVED5 (DRM_COMMAND_BASE + 4)
#define DRM_PVR_RESERVED6 (DRM_COMMAND_BASE + 5)

/* PVR includes */
#define SUPPORT_MEMINFO_IDS
#define SUPPORT_DRI_DRM_EXT
#include <config_kernel.h>
#include <sgxfeaturedefs.h>
#include <pvr_bridge.h>
#include <sgx_bridge.h>
#include <pvr_drm_shared.h>

#define DRM_IOCTL_PVR_SRVKM DRM_IOWR(PVR_DRM_SRVKM_CMD, PVRSRV_BRIDGE_PACKAGE)

#include <pprint.c>

#define UNUSED(x) (void)x

struct mem_alloc {
    PVRSRV_BRIDGE_OUT_ALLOCDEVICEMEM dev_mem;
    void *data;
};

struct driver_state {
    int fd;
    IMG_HANDLE kernel_services;
    IMG_UINT32 device_idx;
    IMG_HANDLE dev_cookie;
    IMG_HANDLE dev_mem_context;
    IMG_HANDLE global_eventobject;

    PVRSRV_HEAP_INFO kernel_heap;
    PVRSRV_HEAP_INFO perctx_3d_heap;
    PVRSRV_HEAP_INFO ta_data_heap;
    struct mem_alloc ccb_buffer;
    struct mem_alloc kernel_buffer_2;
    struct mem_alloc kernel_buffer_3;
    struct mem_alloc perctx_3d_buffer;
    struct mem_alloc ta_data_buffer;

    IMG_HANDLE hw_render_ctx;
    IMG_DEV_VIRTADDR hw_render_ctx_addr;
};

static void check_pvr_error(PVRSRV_ERROR error) {
    if(error != PVRSRV_OK) {
        printf("Got error %s\n", PPRINT(NULL, error, PVRSRV_ERROR));
        assert(false);
    }
}

static void check_ioctl(int fd, long request, void *data) {
    int ret = drmIoctl(fd, request, data);
    if(ret != 0) {
        printf("IOCTL ret value %d, %d, %s\n", ret, errno, strerror(errno));
        assert(ret == 0);
    }
}

static void connect_service(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_CONNECT_SERVICES in = {
        .ui32BridgeFlags = 0xDEADBEEF
    };
    PVRSRV_BRIDGE_OUT_CONNECT_SERVICES out = {.eError = 0};
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_CONNECT_SERVICES,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = 0,
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    //PPRINT(stdout, &out, PVRSRV_BRIDGE_OUT_CONNECT_SERVICES);
    check_pvr_error(out.eError);
    state->kernel_services = out.hKernelServices;
}

static void enum_devices(struct driver_state *state) {
    PVRSRV_BRIDGE_OUT_ENUMDEVICE out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_ENUM_DEVICES,
        .ui32Size = sizeof(data),
        .pvParamIn = NULL,
        .ui32InBufferSize = 0,
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    bool gpu_found = false;
    for(uint32_t i = 0; i < out.ui32NumDevices; i++) {
        PVRSRV_DEVICE_IDENTIFIER *di = &out.asDeviceIdentifier[i];
        if(di->eDeviceType == PVRSRV_DEVICE_TYPE_SGX && di->eDeviceClass == PVRSRV_DEVICE_CLASS_3D) {
            state->device_idx = di->ui32DeviceIndex;
            gpu_found = true;
            break;
        }
    }

    assert(gpu_found);
}

static void acquire_devinfo(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_ACQUIRE_DEVICEINFO in = {
        .uiDevIndex = state->device_idx,
        .eDeviceType = PVRSRV_DEVICE_TYPE_UNKNOWN,
    };
    PVRSRV_BRIDGE_OUT_ACQUIRE_DEVICEINFO out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_ACQUIRE_DEVICEINFO,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    state->dev_cookie = out.hDevCookie;
}

static void sgx_get_misc_info(struct driver_state *state) {
    SGX_MISC_INFO misc_info = {
        .eRequest = SGX_MISC_INFO_REQUEST_DRIVER_SGXREV,
    };
    PVRSRV_BRIDGE_IN_SGXGETMISCINFO in = {
        .hDevCookie = state->dev_cookie,
        .psMiscInfo = &misc_info,
    };
    PVRSRV_BRIDGE_RETURN out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_SGX_GETMISCINFO,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services,
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    /* TODO figure out what we do with misc info */
    PPRINT(stdout, &in, PVRSRV_BRIDGE_IN_SGXGETMISCINFO);
    fprintf(stdout, "\n");
    PPRINT(stdout, &in.psMiscInfo->uData.sSGXFeatures, PVRSRV_SGX_MISCINFO_FEATURES);
    fprintf(stdout, "\n");
}

static void create_devmemcontext(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_CREATE_DEVMEMCONTEXT in = {
        .hDevCookie = state->dev_cookie
    };
    PVRSRV_BRIDGE_OUT_CREATE_DEVMEMCONTEXT out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_CREATE_DEVMEMCONTEXT,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);
    PPRINT(stdout, &out, PVRSRV_BRIDGE_OUT_CREATE_DEVMEMCONTEXT);

    state->dev_mem_context = out.hDevMemContext;
}

static void get_clientinfo(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_GETCLIENTINFO in = {
        .hDevCookie = state->dev_cookie,
    };
    PVRSRV_BRIDGE_OUT_GETCLIENTINFO out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_SGX_GETCLIENTINFO,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);
    //PPRINT(stdout, &out, PVRSRV_BRIDGE_OUT_GETCLIENTINFO);
}

static void get_miscinfo(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_GET_MISC_INFO in = {
        .sMiscInfo = {
            .ui32StateRequest = PVRSRV_MISC_INFO_GLOBALEVENTOBJECT_PRESENT | PVRSRV_MISC_INFO_TIMER_PRESENT,
        },
    };
    PVRSRV_BRIDGE_OUT_GET_MISC_INFO out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_GET_MISC_INFO,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    state->global_eventobject = out.sMiscInfo.sGlobalEventObject.hOSEventKM;
}

static void open_eventobject(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_EVENT_OBJECT_OPEN in = {
        .sEventObject = {
            .szName = "PVRSRV_GLOBAL_EVENTOBJECT",
            .hOSEventKM = state->global_eventobject,
        }
    };

    PVRSRV_BRIDGE_OUT_EVENT_OBJECT_OPEN out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_EVENT_OBJECT_OPEN,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };
    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);
}

static void get_dev_mem_heapinfo(struct driver_state *state) {
    PVRSRV_BRIDGE_IN_GET_DEVMEM_HEAPINFO in = {
        .hDevCookie = state->dev_cookie,
        .hDevMemContext = state->dev_mem_context,
    };
    PVRSRV_BRIDGE_OUT_GET_DEVMEM_HEAPINFO out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_GET_DEVMEM_HEAPINFO,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    for (int i = 0; i < out.ui32ClientHeapCount; i++) {
        //PPRINT(stdout, &out.sHeapInfo[i], PVRSRV_HEAP_INFO);
        if (out.sHeapInfo[i].ui32HeapID == HEAP_ID( PVRSRV_DEVICE_TYPE_SGX, SGX_KERNEL_DATA_HEAP_ID)) {
            state->kernel_heap = out.sHeapInfo[i];
        }

        if (out.sHeapInfo[i].ui32HeapID == HEAP_ID( PVRSRV_DEVICE_TYPE_SGX, SGX_TADATA_HEAP_ID)) {
            state->ta_data_heap = out.sHeapInfo[i];
        }

        if (out.sHeapInfo[i].ui32HeapID == HEAP_ID( PVRSRV_DEVICE_TYPE_SGX, SGX_PERCONTEXT_3DPARAMETERS_HEAP_ID)) {
            state->perctx_3d_heap = out.sHeapInfo[i];
        }
    }

    assert(state->kernel_heap.ui32HeapID != 0);
    assert(state->ta_data_heap.ui32HeapID != 0);
    assert(state->perctx_3d_heap.ui32HeapID != 0);
}

static struct mem_alloc allocate_memobj(struct driver_state *state, PVRSRV_BRIDGE_IN_ALLOCDEVICEMEM *alloc) {
    struct mem_alloc mem;
    PVRSRV_BRIDGE_OUT_ALLOCDEVICEMEM out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_ALLOC_DEVICEMEM,
        .ui32Size = sizeof(data),
        .pvParamIn = alloc,
        .ui32InBufferSize = sizeof(*alloc),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    printf("Allocating from %p %p 0x%x\n", alloc->hDevCookie, alloc->hDevMemHeap, alloc->uSize);
    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    {
        PVRSRV_BRIDGE_IN_MHANDLE_TO_MMAP_DATA in = {
            .hMHandle = out.sClientMemInfo.hKernelMemInfo,
        };
        PVRSRV_BRIDGE_OUT_MHANDLE_TO_MMAP_DATA out2;
        PVRSRV_BRIDGE_PACKAGE data = {
            .ui32BridgeID = PVRSRV_BRIDGE_MHANDLE_TO_MMAP_DATA,
            .ui32Size = sizeof(data),
            .pvParamIn = &in,
            .ui32InBufferSize = sizeof(in),
            .pvParamOut = &out2,
            .ui32OutBufferSize = sizeof(out2),
            .hKernelServices = state->kernel_services
        };

        check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
        check_pvr_error(out2.eError);

        PPRINT(stdout, &out2, PVRSRV_BRIDGE_OUT_MHANDLE_TO_MMAP_DATA);

        printf("Attempting map 0x%x, 0x%x\n", (uint32_t)out2.uiRealByteSize, (uint32_t)out2.uiMMapOffset);
        mem.data = (void*)(uintptr_t)syscall(SYS_mmap2, NULL, out2.uiRealByteSize,
                                             PROT_READ | PROT_WRITE, MAP_SHARED,
                                             state->fd, out2.uiMMapOffset);

        mem.data += out2.uiByteOffset;
        assert(mem.data != MAP_FAILED);
        memset(mem.data, 0x00, out.sClientMemInfo.uAllocSize);
    }

    mem.dev_mem = out;
    return mem;
}

static void allocate_memory(struct driver_state *state) {
     PVRSRV_BRIDGE_IN_ALLOCDEVICEMEM alloc_CCB = {
         .hDevCookie = state->dev_cookie,
         .hDevMemHeap = state->kernel_heap.hDevMemHeap,
         .ui32Attribs = PVRSRV_MEM_READ | PVRSRV_MEM_WRITE | PVRSRV_MEM_CACHE_CONSISTENT | PVRSRV_MEM_NO_SYNCOBJ | PVRSRV_MEM_EDM_PROTECT,
         .uSize = 0x86a4,
         .uAlignment = 0x1000,
     };
     state->ccb_buffer = allocate_memobj(state, &alloc_CCB);

     PVRSRV_BRIDGE_IN_ALLOCDEVICEMEM alloc_kernel_buffer_2 = {
         .hDevCookie = state->dev_cookie,
         .hDevMemHeap = state->kernel_heap.hDevMemHeap,
         .ui32Attribs = PVRSRV_MEM_READ | PVRSRV_MEM_WRITE | PVRSRV_MEM_CACHE_CONSISTENT | PVRSRV_MEM_NO_SYNCOBJ | PVRSRV_MEM_EDM_PROTECT,
         .uSize = 0x8,
         .uAlignment = 0x20,
     };
     state->kernel_buffer_2 = allocate_memobj(state, &alloc_kernel_buffer_2);

     PVRSRV_BRIDGE_IN_ALLOCDEVICEMEM alloc_kernel_buffer_3 = {
         .hDevCookie = state->dev_cookie,
         .hDevMemHeap = state->kernel_heap.hDevMemHeap,
         .ui32Attribs = PVRSRV_MEM_READ | PVRSRV_MEM_WRITE | PVRSRV_MEM_CACHE_CONSISTENT | PVRSRV_MEM_NO_SYNCOBJ,
         .uSize = 0x44,
         .uAlignment = 0x20,
     };
     state->kernel_buffer_3 = allocate_memobj(state, &alloc_kernel_buffer_3);

     PVRSRV_BRIDGE_IN_ALLOCDEVICEMEM alloc_perctx_3d_buffer = {
         .hDevCookie = state->dev_cookie,
         .hDevMemHeap = state->perctx_3d_heap.hDevMemHeap,
         .ui32Attribs = PVRSRV_MEM_READ | PVRSRV_MEM_WRITE | PVRSRV_MEM_NO_SYNCOBJ,
         .uSize = 0x40,
         .uAlignment = 0x40,
     };
     state->perctx_3d_buffer = allocate_memobj(state, &alloc_perctx_3d_buffer);

     PVRSRV_BRIDGE_IN_ALLOCDEVICEMEM alloc_ta_data_buffer = {
         .hDevCookie = state->dev_cookie,
         .hDevMemHeap = state->ta_data_heap.hDevMemHeap,
         .ui32Attribs = PVRSRV_MEM_READ | PVRSRV_MEM_WRITE | PVRSRV_MEM_NO_SYNCOBJ,
         .uSize = 0xc0,
         .uAlignment = 0x40,
     };
     state->ta_data_buffer = allocate_memobj(state, &alloc_ta_data_buffer);
}

static void register_hw_render_ctx(struct driver_state *state) {
    /* Data copied directly from dump on hw
     * Not sure what other params mean, but the buffer
     * ones are obvious once you start offsetting buffer addresses
     */
    uint32_t render_ctx[] = {
        0x21,
        0x1,
        0x0,
        state->ccb_buffer.dev_mem.sClientMemInfo.sDevVAddr.uiAddr, // Needs to be patched
        state->kernel_buffer_2.dev_mem.sClientMemInfo.sDevVAddr.uiAddr, // All zeros
        0x0,
        state->kernel_buffer_3.dev_mem.sClientMemInfo.sDevVAddr.uiAddr, // All zeros
        state->perctx_3d_buffer.dev_mem.sClientMemInfo.sDevVAddr.uiAddr, // All zeros
        state->ta_data_buffer.dev_mem.sClientMemInfo.sDevVAddr.uiAddr, // All zeros
        0x0,
        0x0,
        0x0,
        0x0,
        0x0,
        0x0,
        0x0,
        0x0,
        0x0, // This value seems to change when dumping, but overriding to zero doesn't break anything
    };
    PVRSRV_BRIDGE_IN_SGX_REGISTER_HW_RENDER_CONTEXT in = {
        .hDevCookie = state->dev_cookie,
        .pHWRenderContextCpuVAddr = render_ctx,
        .ui32HWRenderContextSize = sizeof(render_ctx),
        .hDevMemContext = state->dev_mem_context,
    };
    PVRSRV_BRIDGE_OUT_SGX_REGISTER_HW_RENDER_CONTEXT out;
    PVRSRV_BRIDGE_PACKAGE data = {
        .ui32BridgeID = PVRSRV_BRIDGE_SGX_REGISTER_HW_RENDER_CONTEXT,
        .ui32Size = sizeof(data),
        .pvParamIn = &in,
        .ui32InBufferSize = sizeof(in),
        .pvParamOut = &out,
        .ui32OutBufferSize = sizeof(out),
        .hKernelServices = state->kernel_services
    };

    check_ioctl(state->fd, DRM_IOCTL_PVR_SRVKM, &data);
    check_pvr_error(out.eError);

    state->hw_render_ctx = out.hHWRenderContext;
    state->hw_render_ctx_addr = out.sHWRenderContextDevVAddr;
}

int main(int argc, char *argv[]) {
    UNUSED(argc);
    UNUSED(argv);
    printf("Hello world!\n");
    struct driver_state state;

    state.fd = open("/dev/dri/renderD128", O_RDWR);
    assert(state.fd != -1 && "Failed to device file");

    drmVersionPtr version = drmGetVersion(state.fd);
    assert(version && "drmGetVersion failed");

    printf("Driver: %s\n", version->name);
    printf("Version %d.%d.%d\n", version->version_major, version->version_minor, version->version_patchlevel);
    printf("Date: %s\n", version->date);
    printf("Desc: %s\n", version->desc);
    printf("Uniq: %s\n", drmGetBusid(state.fd));

    //printf("DRM_IOCTL_VERSION = 0x%x\n", DRM_IOCTL_VERSION);
    //printf("PVR_DRM_SRVKM_CMD = 0x%x\n", DRM_IOCTL_PVR_SRVKM);
    connect_service(&state);
    enum_devices(&state);
    acquire_devinfo(&state);
    sgx_get_misc_info(&state);
    create_devmemcontext(&state);
    get_clientinfo(&state);
    get_miscinfo(&state);
    open_eventobject(&state);
    get_dev_mem_heapinfo(&state);
    allocate_memory(&state);
    register_hw_render_ctx(&state);

    close(state.fd);

    return 0;
}