// Copyright 2014 Citra Emulator Project
// Licensed under GPLv2 or any later version
// Refer to the license.txt file included.

#include "common/alignment.h"
#include "common/common_types.h"
#include "common/logging/log.h"
#include "core/arm/arm_interface.h"
#include "core/core.h"
#include "core/hle/ipc_helpers.h"
#include "core/hle/kernel/process.h"
#include "core/hle/service/ldr_ro/cro_helper.h"
#include "core/hle/service/ldr_ro/ldr_ro.h"

namespace Service::LDR {

static const ResultCode ERROR_ALREADY_INITIALIZED = // 0xD9612FF9
    ResultCode(ErrorDescription::AlreadyInitialized, ErrorModule::RO, ErrorSummary::Internal,
               ErrorLevel::Permanent);
static const ResultCode ERROR_NOT_INITIALIZED = // 0xD9612FF8
    ResultCode(ErrorDescription::NotInitialized, ErrorModule::RO, ErrorSummary::Internal,
               ErrorLevel::Permanent);
static const ResultCode ERROR_BUFFER_TOO_SMALL = // 0xE0E12C1F
    ResultCode(static_cast<ErrorDescription>(31), ErrorModule::RO, ErrorSummary::InvalidArgument,
               ErrorLevel::Usage);
static const ResultCode ERROR_MISALIGNED_ADDRESS = // 0xD9012FF1
    ResultCode(ErrorDescription::MisalignedAddress, ErrorModule::RO, ErrorSummary::WrongArgument,
               ErrorLevel::Permanent);
static const ResultCode ERROR_MISALIGNED_SIZE = // 0xD9012FF2
    ResultCode(ErrorDescription::MisalignedSize, ErrorModule::RO, ErrorSummary::WrongArgument,
               ErrorLevel::Permanent);
static const ResultCode ERROR_ILLEGAL_ADDRESS = // 0xE1612C0F
    ResultCode(static_cast<ErrorDescription>(15), ErrorModule::RO, ErrorSummary::Internal,
               ErrorLevel::Usage);
static const ResultCode ERROR_INVALID_MEMORY_STATE = // 0xD8A12C08
    ResultCode(static_cast<ErrorDescription>(8), ErrorModule::RO, ErrorSummary::InvalidState,
               ErrorLevel::Permanent);
static const ResultCode ERROR_NOT_LOADED = // 0xD8A12C0D
    ResultCode(static_cast<ErrorDescription>(13), ErrorModule::RO, ErrorSummary::InvalidState,
               ErrorLevel::Permanent);

static bool VerifyBufferState(Kernel::Process& process, VAddr buffer_ptr, u32 size) {
    auto vma = process.vm_manager.FindVMA(buffer_ptr);
    while (vma != process.vm_manager.vma_map.end()) {
        if (vma->second.permissions != Kernel::VMAPermission::ReadWrite ||
            vma->second.meminfo_state != Kernel::MemoryState::Private) {
            return false;
        }
        if (vma->second.base + vma->second.size >= buffer_ptr + size) {
            return true;
        }
        vma = std::next(vma);
    }
    return false;
}

static void InvalidateCacheRange(Core::System& system, std::vector<std::pair<u32, u32>>& ranges) {
    if (ranges.empty()) {
        return;
    }

    std::sort(ranges.begin(), ranges.end(), [](const std::pair<u32, u32>& x, const std::pair<u32, u32>& y)->bool {
        return x.first < y.first;
    });

    u32 start_address = ranges[0].first;
    u32 length = ranges[0].second;

    for(u32 i = 1; i < ranges.size(); ++i) {
        u32 addr = ranges[i].first;
        u32 size = ranges[i].second;
        u32 offset = addr - start_address;
        if (offset > length + Memory::PAGE_SIZE) {
            system.InvalidateCacheRange(start_address, length);
            //
            start_address = addr;
            length = size;
        } else {
            length = std::max(addr + size, start_address + length) - start_address;
        }
    }

    system.InvalidateCacheRange(start_address, length);
    ranges.clear();
}

void RO::Initialize(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x01, 3, 2);
    VAddr crs_buffer_ptr = rp.Pop<u32>();
    u32 crs_size = rp.Pop<u32>();
    VAddr crs_address = rp.Pop<u32>();
    // TODO (wwylele): RO service checks the descriptor here and return error 0xD9001830 for
    // incorrect descriptor. This error return should be probably built in IPC::RequestParser.
    // All other service functions below have the same issue.
    auto process = rp.PopObject<Kernel::Process>();

    LOG_DEBUG(Service_LDR, "called, crs_buffer_ptr=0x{:08X}, crs_address=0x{:08X}, crs_size=0x{:X}",
              crs_buffer_ptr, crs_address, crs_size);

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);

    ClientSlot* slot = GetSessionData(ctx.Session());
    if (slot->loaded_crs != 0) {
        LOG_ERROR(Service_LDR, "Already initialized");
        rb.Push(ERROR_ALREADY_INITIALIZED);
        return;
    }

    if (crs_size < CRO_HEADER_SIZE) {
        LOG_ERROR(Service_LDR, "CRS is too small");
        rb.Push(ERROR_BUFFER_TOO_SMALL);
        return;
    }

    if (crs_buffer_ptr & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRS original address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        return;
    }

    if (crs_address & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRS mapping address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        return;
    }

    if (crs_size & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRS size is not aligned");
        rb.Push(ERROR_MISALIGNED_SIZE);
        return;
    }

    if (!VerifyBufferState(*process, crs_buffer_ptr, crs_size)) {
        LOG_ERROR(Service_LDR, "CRS original buffer is in invalid state");
        rb.Push(ERROR_INVALID_MEMORY_STATE);
        return;
    }

    if (crs_address < Memory::PROCESS_IMAGE_VADDR ||
        crs_address + crs_size > Memory::PROCESS_IMAGE_VADDR_END) {
        LOG_ERROR(Service_LDR, "CRS mapping address is not in the process image region");
        rb.Push(ERROR_ILLEGAL_ADDRESS);
        return;
    }

    ResultCode result = RESULT_SUCCESS;

    result = process->Map(crs_address, crs_buffer_ptr, crs_size, Kernel::VMAPermission::Read, true);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error mapping memory block {:08X}", result.raw);
        rb.Push(result);
        return;
    }

    CROHelper crs(crs_address, *process, system);
    crs.InitCRS();

    result = crs.Rebase(0, crs_size, 0, 0, 0, 0, true);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error rebasing CRS 0x{:08X}", result.raw);
        rb.Push(result);
        return;
    }

    slot->loaded_crs = crs_address;

    InvalidateCacheRange(system, crs.invalidate_cache_ranges);

    rb.Push(RESULT_SUCCESS);
}

void RO::LoadCRR(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x02, 2, 2);
    VAddr crr_buffer_ptr = rp.Pop<u32>();
    u32 crr_size = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);
    rb.Push(RESULT_SUCCESS);

    LOG_WARNING(Service_LDR, "(STUBBED) called, crr_buffer_ptr=0x{:08X}, crr_size=0x{:08X}",
                crr_buffer_ptr, crr_size);
}

void RO::UnloadCRR(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x03, 1, 2);
    u32 crr_buffer_ptr = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);
    rb.Push(RESULT_SUCCESS);

    LOG_WARNING(Service_LDR, "(STUBBED) called, crr_buffer_ptr=0x{:08X}", crr_buffer_ptr);
}

void RO::LoadCRO(Kernel::HLERequestContext& ctx, bool link_on_load_bug_fix) {
    IPC::RequestParser rp(ctx, link_on_load_bug_fix ? 0x09 : 0x04, 11, 2);
    VAddr cro_buffer_ptr = rp.Pop<u32>();
    VAddr cro_address = rp.Pop<u32>();
    u32 cro_size = rp.Pop<u32>();
    VAddr data_segment_address = rp.Pop<u32>();
    u32 zero = rp.Pop<u32>();
    u32 data_segment_size = rp.Pop<u32>();
    u32 bss_segment_address = rp.Pop<u32>();
    u32 bss_segment_size = rp.Pop<u32>();
    bool auto_link = rp.Pop<bool>();
    u32 fix_level = rp.Pop<u32>();
    VAddr crr_address = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    LOG_DEBUG(Service_LDR,
              "LoadCRO called ({}), cro_buffer_ptr=0x{:08X}, cro_address=0x{:08X}, cro_size=0x{:X}, "
              "data_segment_address=0x{:08X}, zero={}, data_segment_size=0x{:X}, "
              "bss_segment_address=0x{:08X}, bss_segment_size=0x{:X}, auto_link={}, "
              "fix_level={}, crr_address=0x{:08X}",
              link_on_load_bug_fix ? "new" : "old", cro_buffer_ptr, cro_address, cro_size,
              data_segment_address, zero, data_segment_size, bss_segment_address, bss_segment_size,
              auto_link ? "true" : "false", fix_level, crr_address);

    IPC::RequestBuilder rb = rp.MakeBuilder(2, 0);

    ClientSlot* slot = GetSessionData(ctx.Session());
    if (slot->loaded_crs == 0) {
        LOG_ERROR(Service_LDR, "Not initialized");
        rb.Push(ERROR_NOT_INITIALIZED);
        rb.Push<u32>(0);
        return;
    }

    if (cro_size < CRO_HEADER_SIZE) {
        LOG_ERROR(Service_LDR, "CRO too small");
        rb.Push(ERROR_BUFFER_TOO_SMALL);
        rb.Push<u32>(0);
        return;
    }

    if (cro_buffer_ptr & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRO original address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        rb.Push<u32>(0);
        return;
    }

    if (cro_address & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRO mapping address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        rb.Push<u32>(0);
        return;
    }

    if (cro_size & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRO size is not aligned");
        rb.Push(ERROR_MISALIGNED_SIZE);
        rb.Push<u32>(0);
        return;
    }

    if (!VerifyBufferState(*process, cro_buffer_ptr, cro_size)) {
        LOG_ERROR(Service_LDR, "CRO original buffer is in invalid state");
        rb.Push(ERROR_INVALID_MEMORY_STATE);
        rb.Push<u32>(0);
        return;
    }

    if (cro_address < Memory::PROCESS_IMAGE_VADDR ||
        cro_address + cro_size > Memory::PROCESS_IMAGE_VADDR_END) {
        LOG_ERROR(Service_LDR, "CRO mapping address is not in the process image region");
        rb.Push(ERROR_ILLEGAL_ADDRESS);
        rb.Push<u32>(0);
        return;
    }

    if (zero) {
        LOG_ERROR(Service_LDR, "Zero is not zero {}", zero);
        rb.Push(ResultCode(static_cast<ErrorDescription>(29), ErrorModule::RO,
                           ErrorSummary::Internal, ErrorLevel::Usage));
        rb.Push<u32>(0);
        return;
    }

    ResultCode result = RESULT_SUCCESS;

    result = process->Map(cro_address, cro_buffer_ptr, cro_size, Kernel::VMAPermission::Read, true);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error mapping memory block {:08X}", result.raw);
        rb.Push(result);
        rb.Push<u32>(0);
        return;
    }

    CROHelper cro(cro_address, *process, system);

    result = cro.VerifyHash(cro_size, crr_address);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error verifying CRO in CRR {:08X}", result.raw);
        process->Unmap(cro_address, cro_buffer_ptr, cro_size, Kernel::VMAPermission::ReadWrite,
                       true);
        rb.Push(result);
        rb.Push<u32>(0);
        return;
    }

    result = cro.Rebase(slot->loaded_crs, cro_size, data_segment_address, data_segment_size,
                        bss_segment_address, bss_segment_size, false);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error rebasing CRO {:08X}", result.raw);
        process->Unmap(cro_address, cro_buffer_ptr, cro_size, Kernel::VMAPermission::ReadWrite,
                       true);
        rb.Push(result);
        rb.Push<u32>(0);
        return;
    }

    result = cro.Link(slot->loaded_crs, link_on_load_bug_fix);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error linking CRO {:08X}", result.raw);
        process->Unmap(cro_address, cro_buffer_ptr, cro_size, Kernel::VMAPermission::ReadWrite,
                       true);
        rb.Push(result);
        rb.Push<u32>(0);
        return;
    }

    cro.Register(slot->loaded_crs, auto_link);

    u32 fix_size = cro.Fix(fix_level);

    if (fix_size != cro_size) {
        result = process->Unmap(cro_address + fix_size, cro_buffer_ptr + fix_size,
                                cro_size - fix_size, Kernel::VMAPermission::ReadWrite, true);
        if (result.IsError()) {
            LOG_ERROR(Service_LDR, "Error unmapping memory block {:08X}", result.raw);
            process->Unmap(cro_address, cro_buffer_ptr, cro_size, Kernel::VMAPermission::ReadWrite,
                           true);
            rb.Push(result);
            rb.Push<u32>(0);
            return;
        }
    }

    auto [exe_begin, exe_size] = cro.GetExecutablePages();
    if (exe_begin) {
        result = process->vm_manager.ReprotectRange(exe_begin, exe_size,
                                                    Kernel::VMAPermission::ReadExecute);
        if (result.IsError()) {
            LOG_ERROR(Service_LDR, "Error reprotecting memory block {:08X}", result.raw);
            process->Unmap(cro_address, cro_buffer_ptr, cro_size, Kernel::VMAPermission::ReadWrite,
                           true);
            rb.Push(result);
            rb.Push<u32>(0);
            return;
        }
    }

    cro.invalidate_cache_ranges.emplace_back(cro_address, cro_size);
    InvalidateCacheRange(system, cro.invalidate_cache_ranges);

    LOG_INFO(Service_LDR, "CRO \"{}\" loaded at 0x{:08X}, fixed_end=0x{:08X}", cro.ModuleName(),
             cro_address, cro_address + fix_size);

    rb.Push(RESULT_SUCCESS, fix_size);
}

void RO::UnloadCRO(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x05, 3, 2);
    VAddr cro_address = rp.Pop<u32>();
    u32 zero = rp.Pop<u32>();
    VAddr cro_buffer_ptr = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    CROHelper cro(cro_address, *process, system);

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);

    ClientSlot* slot = GetSessionData(ctx.Session());
    if (slot->loaded_crs == 0) {
        LOG_ERROR(Service_LDR, "Not initialized");
        rb.Push(ERROR_NOT_INITIALIZED);
        return;
    }

    if (cro_address & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRO address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        return;
    }

    if (!cro.IsLoaded()) {
        LOG_ERROR(Service_LDR, "Invalid or not loaded CRO");
        rb.Push(ERROR_NOT_LOADED);
        return;
    }

    LOG_INFO(Service_LDR, "UnloadCRO CRO \"{}\", cro_address=0x{:08X}, zero={}, cro_buffer_ptr=0x{:08X}",
             cro.ModuleName(), cro_address, zero, cro_buffer_ptr);

    u32 fixed_size = cro.GetFixedSize();

    cro.Unregister(slot->loaded_crs);

    ResultCode result = cro.Unlink(slot->loaded_crs);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error unlinking CRO {:08X}", result.raw);
        rb.Push(result);
        return;
    }

    // If the module is not fixed, clears all external/internal relocations
    // to restore the state before loading, so that it can be loaded again(?)
    if (!cro.IsFixed()) {
        result = cro.ClearRelocations();
        if (result.IsError()) {
            LOG_ERROR(Service_LDR, "Error clearing relocations {:08X}", result.raw);
            rb.Push(result);
            return;
        }
    }

    cro.Unrebase(false);

    result = process->Unmap(cro_address, cro_buffer_ptr, fixed_size,
                            Kernel::VMAPermission::ReadWrite, true);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error unmapping CRO {:08X}", result.raw);
    }

    cro.invalidate_cache_ranges.emplace_back(cro_address, fixed_size);
    InvalidateCacheRange(system, cro.invalidate_cache_ranges);

    rb.Push(result);
}

void RO::LinkCRO(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x06, 1, 2);
    VAddr cro_address = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    LOG_DEBUG(Service_LDR, "LinkCRO called, cro_address=0x{:08X}", cro_address);

    CROHelper cro(cro_address, *process, system);

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);

    ClientSlot* slot = GetSessionData(ctx.Session());
    if (slot->loaded_crs == 0) {
        LOG_ERROR(Service_LDR, "Not initialized");
        rb.Push(ERROR_NOT_INITIALIZED);
        return;
    }

    if (cro_address & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRO address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        return;
    }

    if (!cro.IsLoaded()) {
        LOG_ERROR(Service_LDR, "Invalid or not loaded CRO");
        rb.Push(ERROR_NOT_LOADED);
        return;
    }

    LOG_INFO(Service_LDR, "Linking CRO \"{}\"", cro.ModuleName());

    ResultCode result = cro.Link(slot->loaded_crs, false);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error linking CRO {:08X}", result.raw);
    }

    InvalidateCacheRange(system, cro.invalidate_cache_ranges);

    rb.Push(result);
}

void RO::UnlinkCRO(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x07, 1, 2);
    VAddr cro_address = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    LOG_DEBUG(Service_LDR, "called, cro_address=0x{:08X}", cro_address);

    CROHelper cro(cro_address, *process, system);

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);

    ClientSlot* slot = GetSessionData(ctx.Session());
    if (slot->loaded_crs == 0) {
        LOG_ERROR(Service_LDR, "Not initialized");
        rb.Push(ERROR_NOT_INITIALIZED);
        return;
    }

    if (cro_address & Memory::PAGE_MASK) {
        LOG_ERROR(Service_LDR, "CRO address is not aligned");
        rb.Push(ERROR_MISALIGNED_ADDRESS);
        return;
    }

    if (!cro.IsLoaded()) {
        LOG_ERROR(Service_LDR, "Invalid or not loaded CRO");
        rb.Push(ERROR_NOT_LOADED);
        return;
    }

    LOG_INFO(Service_LDR, "Unlinking CRO \"{}\"", cro.ModuleName());

    ResultCode result = cro.Unlink(slot->loaded_crs);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error unlinking CRO {:08X}", result.raw);
    }

    InvalidateCacheRange(system, cro.invalidate_cache_ranges);

    rb.Push(result);
}

void RO::Shutdown(Kernel::HLERequestContext& ctx) {
    IPC::RequestParser rp(ctx, 0x08, 1, 2);
    VAddr crs_buffer_ptr = rp.Pop<u32>();
    auto process = rp.PopObject<Kernel::Process>();

    LOG_DEBUG(Service_LDR, "called, crs_buffer_ptr=0x{:08X}", crs_buffer_ptr);

    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);

    ClientSlot* slot = GetSessionData(ctx.Session());
    if (slot->loaded_crs == 0) {
        LOG_ERROR(Service_LDR, "Not initialized");
        rb.Push(ERROR_NOT_INITIALIZED);
        return;
    }

    CROHelper crs(slot->loaded_crs, *process, system);
    crs.Unrebase(true);

    ResultCode result = RESULT_SUCCESS;

    result = process->Unmap(slot->loaded_crs, crs_buffer_ptr, crs.GetFileSize(),
                            Kernel::VMAPermission::ReadWrite, true);
    if (result.IsError()) {
        LOG_ERROR(Service_LDR, "Error unmapping CRS {:08X}", result.raw);
    }

    InvalidateCacheRange(system, crs.invalidate_cache_ranges);

    slot->loaded_crs = 0;
    rb.Push(result);
}

RO::RO(Core::System& system) : ServiceFramework("ldr:ro", 2), system(system) {
    static const FunctionInfo functions[] = {
        {0x000100C2, &RO::Initialize, "Initialize"},
        {0x00020082, &RO::LoadCRR, "LoadCRR"},
        {0x00030042, &RO::UnloadCRR, "UnloadCRR"},
        {0x000402C2, &RO::LoadCRO<false>, "LoadCRO"},
        {0x000500C2, &RO::UnloadCRO, "UnloadCRO"},
        {0x00060042, &RO::LinkCRO, "LinkCRO"},
        {0x00070042, &RO::UnlinkCRO, "UnlinkCRO"},
        {0x00080042, &RO::Shutdown, "Shutdown"},
        {0x000902C2, &RO::LoadCRO<true>, "LoadCRO_New"},
    };
    RegisterHandlers(functions);
}

void InstallInterfaces(Core::System& system) {
    auto& service_manager = system.ServiceManager();
    std::make_shared<RO>(system)->InstallAsService(service_manager);
}

} // namespace Service::LDR
