/*
 * Copyright 2012, Alex Smith, alex@alex-smith.me.uk.
 * Copyright 2009-2012, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Copyright 2008, François Revol, revol@free.fr
 * Copyright 2016, Rene Gollent, rene@gollent.com.
 * Distributed under the terms of the MIT License.
 */

#include "DisassemblerX8664.h"

#include <new>

#include "Zycore/Format.h"
#include "Zydis/Zydis.h"

#include <OS.h>


#include "CpuStateX8664.h"
#include "InstructionInfo.h"


void
CpuStateToZydisRegContext(CpuStateX8664* state, ZydisRegisterContext* context)
{
        context->values[ZYDIS_REGISTER_RAX] = state->IntRegisterValue(X86_64_REGISTER_RAX);
        context->values[ZYDIS_REGISTER_RSP] = state->IntRegisterValue(X86_64_REGISTER_RSP);
        context->values[ZYDIS_REGISTER_RIP] = state->IntRegisterValue(X86_64_REGISTER_RIP);
        // context->values[ZYDIS_REGISTER_RFLAGS] = eflags;
        context->values[ZYDIS_REGISTER_RCX] = state->IntRegisterValue(X86_64_REGISTER_RCX);
        context->values[ZYDIS_REGISTER_RDX] = state->IntRegisterValue(X86_64_REGISTER_RDX);
        context->values[ZYDIS_REGISTER_RBX] = state->IntRegisterValue(X86_64_REGISTER_RBX);
        context->values[ZYDIS_REGISTER_RBP] = state->IntRegisterValue(X86_64_REGISTER_RBP);
        context->values[ZYDIS_REGISTER_RSI] = state->IntRegisterValue(X86_64_REGISTER_RSI);
        context->values[ZYDIS_REGISTER_RDI] = state->IntRegisterValue(X86_64_REGISTER_RDI);
        context->values[ZYDIS_REGISTER_R8] = state->IntRegisterValue(X86_64_REGISTER_R8);
        context->values[ZYDIS_REGISTER_R9] = state->IntRegisterValue(X86_64_REGISTER_R9);
        context->values[ZYDIS_REGISTER_R10] = state->IntRegisterValue(X86_64_REGISTER_R10);
        context->values[ZYDIS_REGISTER_R11] = state->IntRegisterValue(X86_64_REGISTER_R11);
        context->values[ZYDIS_REGISTER_R12] = state->IntRegisterValue(X86_64_REGISTER_R12);
        context->values[ZYDIS_REGISTER_R13] = state->IntRegisterValue(X86_64_REGISTER_R13);
        context->values[ZYDIS_REGISTER_R14] = state->IntRegisterValue(X86_64_REGISTER_R14);
        context->values[ZYDIS_REGISTER_R15] = state->IntRegisterValue(X86_64_REGISTER_R15);
}


struct DisassemblerX8664::ZydisData {
        ZydisDecoder decoder ;
        ZydisFormatter formatter;
        ZyanUSize offset;
};


DisassemblerX8664::DisassemblerX8664()
        :
        fAddress(0),
        fCode(NULL),
        fCodeSize(0),
        fZydisData(NULL)
{
}


DisassemblerX8664::~DisassemblerX8664()
{
        delete fZydisData;
}


status_t
DisassemblerX8664::Init(target_addr_t address, const void* code, size_t codeSize)
{
        // unset old data
        delete fZydisData;
        fZydisData = NULL;

        // set new data
        fZydisData = new(std::nothrow) ZydisData;
        if (fZydisData == NULL)
                return B_NO_MEMORY;

        fAddress = address;
        fCode = (const uint8*)code;
        fCodeSize = codeSize;

        // init zydis
        fZydisData->offset = 0;
        ZydisDecoderInit(&fZydisData->decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_STACK_WIDTH_64);
        ZydisFormatterInit(&fZydisData->formatter, ZYDIS_FORMATTER_STYLE_ATT);
        ZydisFormatterSetProperty(&fZydisData->formatter, ZYDIS_FORMATTER_PROP_FORCE_SIZE, ZYAN_TRUE);
        ZydisFormatterSetProperty(&fZydisData->formatter, ZYDIS_FORMATTER_PROP_HEX_UPPERCASE,
                ZYAN_FALSE);
        ZydisFormatterSetProperty(&fZydisData->formatter, ZYDIS_FORMATTER_PROP_ADDR_PADDING_ABSOLUTE,
                ZYDIS_PADDING_DISABLED);
        ZydisFormatterSetProperty(&fZydisData->formatter, ZYDIS_FORMATTER_PROP_ADDR_PADDING_RELATIVE,
                ZYDIS_PADDING_DISABLED);
        ZydisFormatterSetProperty(&fZydisData->formatter, ZYDIS_FORMATTER_PROP_DISP_PADDING,
                ZYDIS_PADDING_DISABLED);
        ZydisFormatterSetProperty(&fZydisData->formatter, ZYDIS_FORMATTER_PROP_IMM_PADDING,
                ZYDIS_PADDING_DISABLED);
                // TODO: Set the correct vendor!

        return B_OK;
}


status_t
DisassemblerX8664::GetNextInstruction(BString& line, target_addr_t& _address,
        target_size_t& _size, bool& _breakpointAllowed)
{
        const uint8* buffer = fCode + fZydisData->offset;
        ZydisDecodedInstruction instruction;
        ZydisDecodedOperand operands[ZYDIS_MAX_OPERAND_COUNT];
        if (!ZYAN_SUCCESS(ZydisDecoderDecodeFull(&fZydisData->decoder, buffer,
                fCodeSize - fZydisData->offset, &instruction, operands))) {
                return B_ENTRY_NOT_FOUND;
        }

        target_addr_t address = fAddress + fZydisData->offset;
        fZydisData->offset += instruction.length;

        char hexString[32];
        char* srcHex = hexString;
        for (ZyanUSize i = 0; i < instruction.length; i++) {
                sprintf(srcHex, "%02" PRIx8, buffer[i]);
                srcHex += 2;
        }

        char formatted[1024];
        if (ZYAN_SUCCESS(ZydisFormatterFormatInstruction(&fZydisData->formatter, &instruction, operands,
                instruction.operand_count_visible, formatted, sizeof(formatted), address, NULL))) {
                line.SetToFormat("0x%016" B_PRIx64 ": %16.16s  %s", address, hexString, formatted);
        } else {
                line.SetToFormat("0x%016" B_PRIx64 ": failed-to-format", address);
        }
                // TODO: Resolve symbols!

        _address = address;
        _size = instruction.length;
        _breakpointAllowed = true;
                // TODO: Implement (rep!)!

        return B_OK;
}


status_t
DisassemblerX8664::GetPreviousInstruction(target_addr_t nextAddress,
        target_addr_t& _address, target_size_t& _size)
{
        if (nextAddress < fAddress || nextAddress > fAddress + fCodeSize)
                return B_BAD_VALUE;

        // loop until hitting the last instruction
        while (true) {
                const uint8* buffer = fCode + fZydisData->offset;
                ZydisDecodedInstruction instruction;
                if (!ZYAN_SUCCESS(ZydisDecoderDecodeInstruction(&fZydisData->decoder,
                                (ZydisDecoderContext*)ZYAN_NULL, buffer, fCodeSize - fZydisData->offset,
                                &instruction))) {
                        return B_ENTRY_NOT_FOUND;
                }

                fZydisData->offset += instruction.length;
                target_addr_t address = fAddress + fZydisData->offset;
                if (address == nextAddress) {
                        _address = address;
                        _size = instruction.length;
                        return B_OK;
                }
        }
}


status_t
DisassemblerX8664::GetNextInstructionInfo(InstructionInfo& _info,
        CpuState* state)
{
        const uint8* buffer = fCode + fZydisData->offset;
        ZydisDecodedInstruction instruction;
        ZydisDecodedOperand operands[ZYDIS_MAX_OPERAND_COUNT];
        if (!ZYAN_SUCCESS(ZydisDecoderDecodeFull(&fZydisData->decoder, buffer,
                fCodeSize - fZydisData->offset, &instruction, operands))) {
                return B_ENTRY_NOT_FOUND;
        }

        target_addr_t address = fAddress + fZydisData->offset;
        fZydisData->offset += instruction.length;

        char hexString[32];
        char* srcHex = hexString;
        for (ZyanUSize i = 0; i < instruction.length; i++) {
                sprintf(srcHex, "%02" PRIx8, buffer[i]);
                srcHex += 2;
        }

        instruction_type type = INSTRUCTION_TYPE_OTHER;
        target_addr_t targetAddress = 0;
        if (instruction.mnemonic == ZYDIS_MNEMONIC_CALL)
                type = INSTRUCTION_TYPE_SUBROUTINE_CALL;
        else if (instruction.mnemonic == ZYDIS_MNEMONIC_JMP)
                type = INSTRUCTION_TYPE_JUMP;
        if (state != NULL) {
                CpuStateX8664* x64State = dynamic_cast<CpuStateX8664*>(state);
                if (x64State != NULL) {
                        ZydisRegisterContext registers;
                        CpuStateToZydisRegContext(x64State, &registers);
                        ZYAN_CHECK(ZydisCalcAbsoluteAddressEx(&instruction, operands,
                                address, &registers, &targetAddress));
                }
        }

        char string[1024];
        int written = snprintf(string, sizeof(string), "0x%016" B_PRIx64 ": %16.16s  ", address,
                hexString);
        char* formatted = string + written;
        if (!ZYAN_SUCCESS(ZydisFormatterFormatInstruction(&fZydisData->formatter, &instruction,
                operands, instruction.operand_count_visible, formatted, sizeof(string) - written,
                address, NULL))) {
                snprintf(string, sizeof(string), "0x%016" B_PRIx64 ": failed-to-format", address);
        }

                // TODO: Resolve symbols!

        if (!_info.SetTo(address, targetAddress, instruction.length, type, true, string))
                return B_NO_MEMORY;

        return B_OK;
}