#include "function_patcher.h"
#include "FunctionAddressProvider.h"
#include "PatchedFunctionData.h"
#include "utils/CThread.h"
#include "utils/logger.h"
#include <coreinit/cache.h>
#include <coreinit/debug.h>
#include <coreinit/memorymap.h>
#include <kernel/kernel.h>
#include <memory>

static void writeDataAndFlushIC(CThread *thread, void *arg) {
    auto *data = (PatchedFunctionData *) arg;

    uint32_t replace_instruction = data->replaceWithInstruction;
    uint32_t physical_address    = data->realPhysicalFunctionAddress;
    uint32_t effective_address   = data->realEffectiveFunctionAddress;
    DCFlushRange(&replace_instruction, 4);
    DCFlushRange(&physical_address, 4);

    auto replace_instruction_physical = (uint32_t) &replace_instruction;

    if (replace_instruction_physical < 0x00800000 || replace_instruction_physical >= 0x01000000) {
        replace_instruction_physical = OSEffectiveToPhysical(replace_instruction_physical);
    } else {
        replace_instruction_physical = replace_instruction_physical + 0x30800000 - 0x00800000;
    }

    KernelCopyData(physical_address, replace_instruction_physical, 4);
    ICInvalidateRange((void *) (effective_address), 4);
}

bool PatchFunction(std::shared_ptr<PatchedFunctionData> &patchedFunction) {
    if (patchedFunction->isPatched) {
        return true;
    }

    // The addresses of a function might change every time with run another application.
    if (!patchedFunction->updateFunctionAddresses()) {
        return true;
    }

    if (patchedFunction->functionName) {
        DEBUG_FUNCTION_LINE("Patching function %s...", patchedFunction->functionName->c_str());
    } else {
        DEBUG_FUNCTION_LINE("Patching function @ %08X", patchedFunction->realEffectiveFunctionAddress);
    }

    volatile uint32_t replacedInstruction;

    auto targetAddress = (uint32_t) &replacedInstruction;
    if (targetAddress < 0x00800000 || targetAddress >= 0x01000000) {
        targetAddress = (uint32_t) OSEffectiveToPhysical(targetAddress);
    } else {
        targetAddress = targetAddress + 0x30800000 - 0x00800000;
    }

    if (targetAddress == 0) {
        DEBUG_FUNCTION_LINE_ERR("Failed to get physical address");
        OSFatal("Failed to get physical address");
        return false;
    }

    // Save the instruction we will replace.
    KernelCopyData(targetAddress, patchedFunction->realPhysicalFunctionAddress, 4);
    DCFlushRange((void *) &replacedInstruction, 4);
    patchedFunction->replacedInstruction = replacedInstruction;

    // Generate a jump to the original function so the unpatched function can still be called
    patchedFunction->generateJumpToOriginal();

    // Generate a code that is run when somebody calls the patched function.
    // If the correct process calls this, it'll jump the function replacement, otherwise the original function will be called.
    patchedFunction->generateReplacementJump();

    // Write this->replaceWithInstruction to the first instruction of the function we want to replace.
    CThread::runOnAllCores(writeDataAndFlushIC, patchedFunction.get());

    // Set patch status
    patchedFunction->isPatched = true;

    return true;
}

bool RestoreFunction(std::shared_ptr<PatchedFunctionData> &patchedFunction) {
    if (!patchedFunction->isPatched) {
        DEBUG_FUNCTION_LINE_VERBOSE("Skip restoring function because it's not patched");
        return true;
    }
    if (patchedFunction->replacedInstruction == 0 || patchedFunction->realEffectiveFunctionAddress == 0) {
        DEBUG_FUNCTION_LINE_ERR("Failed to restore function, information is missing.");
        return false;
    }

    auto targetAddrPhys = (uint32_t) patchedFunction->realPhysicalFunctionAddress;

    if (patchedFunction->library != LIBRARY_OTHER) {
        targetAddrPhys = (uint32_t) OSEffectiveToPhysical(patchedFunction->realEffectiveFunctionAddress);
    }

    if (patchedFunction->isDynamicFunction() &&
        // Other processes than the wii u menu and game one seem to keep their rpl's loaded.
        patchedFunction->targetProcess != FP_TARGET_PROCESS_GAME_AND_MENU &&
        patchedFunction->targetProcess != FP_TARGET_PROCESS_GAME &&
        patchedFunction->targetProcess != FP_TARGET_PROCESS_WII_U_MENU) {
        DEBUG_FUNCTION_LINE_VERBOSE("Its a dynamic function. We don't need to restore it!");
    } else {
        DEBUG_FUNCTION_LINE_VERBOSE("Restoring %08X to %08X [%08X]", (uint32_t) patchedFunction->replacedInstruction, patchedFunction->realEffectiveFunctionAddress, targetAddrPhys);
        auto sourceAddr = (uint32_t) &patchedFunction->replacedInstruction;

        auto sourceAddrPhys = (uint32_t) OSEffectiveToPhysical(sourceAddr);

        // These hardcoded values should be replaced with something more dynamic.
        if (sourceAddrPhys == 0 && (sourceAddr >= 0x00800000 && sourceAddr < 0x01000000)) {
            sourceAddrPhys = sourceAddr + (0x30800000 - 0x00800000);
        }

        if (sourceAddrPhys == 0) {
            OSFatal("Failed to get physical address");
        }

        KernelCopyData(targetAddrPhys, sourceAddrPhys, 4);
        ICInvalidateRange((void *) patchedFunction->realEffectiveFunctionAddress, 4);
        DCFlushRange((void *) patchedFunction->realEffectiveFunctionAddress, 4);
    }
    patchedFunction->isPatched = false;
    return true;
}