#include <pspkernel.h>
#include <string.h>
#include "libs.h"
#include "misc.h"
#include "nploader.h"
#include "logger.h"

PSP_MODULE_INFO("nploader", PSP_MODULE_KERNEL, 0, 5);
PSP_HEAP_SIZE_KB(2);

STMOD_HANDLER previous = NULL;
int module_found = 0;
int loader_found = 0;

char modname[32];

void patch_io(SceModule2 *module) {
     kprintf("Patching IoFileMgrForUser\n");
    const char * base = "IoFileMgrForUser";
    // sceIoOpen
     sceIoOpen_func = NULL;
    if(hook_import_bynid(module, base, 0x109F50BC, np_open, 1) < 0)
          kprintf(">> hook to sceIoOpen failed\n");
     // sceIoOpenAsync
     sceIoOpenAsync_func = NULL;
    if(hook_import_bynid(module, base, 0x89AA9906, np_open_async, 1) < 0)
          kprintf(">> hook to sceIoOpenAsync failed (not critical)\n");
}

void patch_load(SceModule2 *module) {
     kprintf("Patching ModuleMgrForUser\n");
    const char * base = "ModuleMgrForUser";
    // sceKernelLoadModule2
     sceKernelLoadModule2_func = NULL;
    if(hook_import_bynid(module, base, 0xF2D8D1B4, np_load, 1) < 0)
          kprintf(">> hook to sceKernelLoadModule2 failed (game without sprx support)\n");
}

int patch_drm(SceModule2 *module) {
     kprintf("Patching scePspNpDrm_user\n");
    const char *base = "scePspNpDrm_user";

     sceNpDrmEdataSetupKey_func = NULL;
    // sceNpDrmEdataSetupKey
    if(hook_import_bynid(module, base, 0x08D98894, np_setup, 1) < 0) {
          kprintf(">> hook to sceNpDrmEdataSetupKey failed\n");
          return 0;
     }
     sceNpDrmEdataGetDataSize_func = NULL;
    // sceNpDrmEdataGetDataSize
    if(hook_import_bynid(module, base, 0x219EF5CC, np_size, 1) < 0)
          kprintf(">> hook to sceNpDrmEdataGetDataSize failed (not critical)\n");
     sceNpDrmRenameCheck_func = NULL;
    // sceNpDrmRenameCheck
    if(hook_import_bynid(module, base, 0x275987D1, np_rename, 1) < 0)
          kprintf(">> hook to sceNpDrmRenameCheck failed (not critical)\n");
     return 1;
}

int module_start_handler(SceModule2 * module) {
     kprintf("> Loaded: %s\n", module->modname);
    if(!module_found && module->text_addr == 0x08804000 && module->entry_addr != 0xFFFFFFFF) {
          //blacklist the Prometheus iso loader
          if(!loader_found && (!strcmp(module->modname, "PLoaderGUI"))) {
               kprintf("Prometheus loader found\n");
               loader_found = 1;
          } else {
               kprintf("Game found: %s\n", module->modname);
               if(patch_drm(module)) {
                    patch_io(module);
                    patch_load(module);
               }
               strcpy(modname, module->modname);
               module_found = 1;
          }
    }
     if(module_found && !strcmp(module->modname, "scePspNpDrm_Driver")) {
          SceModule2 *mod = (SceModule2 *)sceKernelFindModuleByName(modname);
          kprintf("> Late scePspNpDrm load, hooking\n");
          patch_drm(mod);
     }
    return previous ? previous(module) : 0;
}

int thread_start(SceSize args, void *argp) {
     kprintf("Making module_handle hook\n");
    previous = sctrlHENSetStartModuleHandler(module_start_handler);
    return 0;
}

int module_start(SceSize argc, void *argp) {
     kprintf("------------------\nNPloader starting\n");
     SceUID thid = sceKernelCreateThread("nploader_main", thread_start, 0x22, 0x2000, 0, NULL);
     if(thid >= 0)
          sceKernelStartThread(thid, argc, argp);
     return 0;
}

int module_stop(SceSize args, void *argp) {
     return 0;
}