/*
 * Copyright 2008, François Revol, revol@free.fr
 * Distributed under the terms of the MIT License.
 */

#include <OS.h>
#include <KernelExport.h>

#include <debug.h>
#include <stdio.h>

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

#include "disasm_arch.h"
#include "elf.h"


static ZydisDecoder sDecoder;
static ZydisFormatter sFormatter;
static ZydisFormatterFunc sDefaultPrintAddressAbsolute;


static ZyanStatus
ZydisFormatterPrintAddressAbsolute(const ZydisFormatter* formatter, ZydisFormatterBuffer* buffer,
        ZydisFormatterContext* context)
{
        ZyanU64 address;
        ZYAN_CHECK(ZydisCalcAbsoluteAddress(context->instruction, context->operand,
                context->runtime_address, &address));

        const char* symbolName;
        addr_t baseAddress;
        status_t error;

        if (IS_KERNEL_ADDRESS(address)) {
                error = elf_debug_lookup_symbol_address(address, &baseAddress,
                        &symbolName, NULL, NULL);
        } else {
                error = elf_debug_lookup_user_symbol_address(
                        debug_get_debugged_thread()->team, address, &baseAddress,
                        &symbolName, NULL, NULL);
        }

        if (error == B_OK) {
                ZYAN_CHECK(ZydisFormatterBufferAppend(buffer, ZYDIS_TOKEN_SYMBOL));
                ZyanString* string;
                ZYAN_CHECK(ZydisFormatterBufferGetString(buffer, &string));
                int64_t offset = address - baseAddress;
                if (offset == 0)
                        return ZyanStringAppendFormat(string, "<%s>", symbolName);
                return ZyanStringAppendFormat(string, "<%s+0x%" B_PRIx64 ">", symbolName, offset);
        }

        return sDefaultPrintAddressAbsolute(formatter, buffer, context);
}


extern "C" void
diasm_arch_assert_fail(const char* assertion, const char* file, unsigned int line,
        const char* function)
{
        kprintf("assert_fail: %s\n", assertion);
        while (true)
                ;
}


extern "C" void
disasm_arch_assert(const char *condition)
{
        kprintf("assert: %s\n", condition);
}


status_t
disasm_arch_dump_insns(addr_t where, int count, addr_t baseAddress,
        int backCount)
{
        ZyanU8 buffer[ZYDIS_MAX_INSTRUCTION_LENGTH];
        ZydisDecodedInstruction instruction;
        int skipCount = 0;

        if (backCount > 0) {
                // count the instructions from base address to start address
                addr_t address = baseAddress;

                int baseCount = 0;

                while (address < where
                        && debug_memcpy(B_CURRENT_TEAM, &buffer, (const void*)address, sizeof(buffer)) == B_OK
                        && ZYAN_SUCCESS(ZydisDecoderDecodeInstruction(&sDecoder,
                                (ZydisDecoderContext*)ZYAN_NULL, buffer, sizeof(buffer), &instruction))) {
                        address += instruction.length;
                        baseCount++;
                }

                if (address == where) {
                        if (baseCount > backCount)
                                skipCount = baseCount - backCount;
                        count += baseCount;
                } else
                        baseAddress = where;
        } else
                baseAddress = where;

        ZyanUSize offset = 0;

        for (int i = 0; i < count; i++, offset += instruction.length) {
                if (debug_memcpy(B_CURRENT_TEAM, &buffer, (const void*)(baseAddress + offset),
                                sizeof(buffer))
                        != B_OK) {
                        kprintf("<read fault>\n");
                        break;
                }
                ZydisDecodedOperand operands[ZYDIS_MAX_OPERAND_COUNT];
                if (!ZYAN_SUCCESS(ZydisDecoderDecodeFull(&sDecoder, buffer, sizeof(buffer), &instruction,
                                operands))) {
                        break;
                }
                if (skipCount > 0) {
                        skipCount--;
                        continue;
                }

                addr_t address = baseAddress + offset;
                if (address == where)
                        kprintf("\x1b[34m");

                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(&sFormatter, &instruction, operands,
                                instruction.operand_count_visible, formatted, sizeof(formatted),
                                baseAddress + offset, NULL))) {
                        kprintf("%#16llx: %16.16s\t%s\n", static_cast<unsigned long long>(address), hexString,
                                formatted);
                } else {
                        kprintf("%#16llx: failed-to-format\n", static_cast<unsigned long long>(address));
                }
                if (address == where)
                        kprintf("\x1b[m");
        }
        return B_OK;
}


status_t
disasm_arch_init()
{
#ifdef __x86_64__
        ZydisDecoderInit(&sDecoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_STACK_WIDTH_64);
#else
        ZydisDecoderInit(&sDecoder, ZYDIS_MACHINE_MODE_LEGACY_32, ZYDIS_STACK_WIDTH_32);
#endif

        ZydisFormatterInit(&sFormatter, ZYDIS_FORMATTER_STYLE_ATT);
        ZydisFormatterSetProperty(&sFormatter, ZYDIS_FORMATTER_PROP_FORCE_SIZE, ZYAN_TRUE);
        ZydisFormatterSetProperty(&sFormatter, ZYDIS_FORMATTER_PROP_HEX_UPPERCASE, ZYAN_FALSE);
        ZydisFormatterSetProperty(&sFormatter, ZYDIS_FORMATTER_PROP_ADDR_PADDING_ABSOLUTE,
                ZYDIS_PADDING_DISABLED);
        ZydisFormatterSetProperty(&sFormatter, ZYDIS_FORMATTER_PROP_ADDR_PADDING_RELATIVE,
                ZYDIS_PADDING_DISABLED);
        ZydisFormatterSetProperty(&sFormatter, ZYDIS_FORMATTER_PROP_DISP_PADDING,
                ZYDIS_PADDING_DISABLED);
        ZydisFormatterSetProperty(&sFormatter, ZYDIS_FORMATTER_PROP_IMM_PADDING,
                ZYDIS_PADDING_DISABLED);

        sDefaultPrintAddressAbsolute = (ZydisFormatterFunc)&ZydisFormatterPrintAddressAbsolute;
        ZydisFormatterSetHook(&sFormatter, ZYDIS_FORMATTER_FUNC_PRINT_ADDRESS_ABS,
                (const void**)&sDefaultPrintAddressAbsolute);

        // XXX: check for AMD and set sVendor;
        return B_OK;
}


status_t
disasm_arch_fini()
{
        return B_OK;
}