#include "disasm/engine.hpp"
#include <cstring>

Disassembler::Disassembler() : m_handle(0), m_initialized(false) {}
Disassembler::~Disassembler()
{
    if (m_initialized) cs_close(&m_handle);
}

bool Disassembler::init(Arch arch)
{
    cs_arch a;
    cs_mode m;
    if (arch == Arch::X64) { a = CS_ARCH_X86; m = CS_MODE_64; }
    else if (arch == Arch::X86) { a = CS_ARCH_X86; m = CS_MODE_32; }
    else return false;

    if (cs_open(a, m, &m_handle) != CS_ERR_OK) return false;
    cs_option(m_handle, CS_OPT_DETAIL, CS_OPT_ON);
    m_initialized = true;
    return true;
}

vector<Instruction> Disassembler::disassemble(const vector<uint8_t>& bytes, addr_t base, size_t count)
{
    vector<Instruction> result;
    if (!m_initialized || bytes.empty()) return result;

    const uint8_t* code = bytes.data();
    size_t size = bytes.size();
    addr_t addr = base;
    size_t n = count ? count : size;

    cs_insn* insn = cs_malloc(m_handle);
    if (!insn) return result;

    while (cs_disasm_iter(m_handle, &code, &size, &addr, insn))
    {
        Instruction inst;
        inst.address = insn->address;
        inst.mnemonic = insn->mnemonic;
        inst.operands = insn->op_str;
        inst.size = static_cast<uint8_t>(insn->size);
        inst.bytes.assign(insn->bytes, insn->bytes + insn->size);
        result.push_back(inst);

        if (count && result.size() >= count) break;
    }

    cs_free(insn, 1);
    return result;
}

bool Disassembler::isProlog(const vector<Instruction>& insts) const
{
    for (const auto& inst : insts)
    {
        if (inst.mnemonic == "push" && (inst.operands == "rbp" || inst.operands == "ebp"))
            return true;
        if (inst.mnemonic == "push")
        {
            const string& ops = inst.operands;
            if (ops == "r15" || ops == "r14" || ops == "r13" || ops == "r12" ||
                ops == "rbx" || ops == "rdi" || ops == "rsi")
                return true;
        }
        if (inst.mnemonic == "sub")
        {
            const string& ops = inst.operands;
            size_t p = ops.find("rsp");
            if (p == string::npos) p = ops.find("esp");
            if (p != string::npos && (p == 0 || ops[p-1] == ' ' || ops[p-1] == ','))
            {
                size_t x = ops.rfind("0x");
                if (x != string::npos && x > p)
                {
                    try {
                        uint64_t val = stoull(ops.substr(x), nullptr, 16);
                        if (val >= 0x20) return true;
                    } catch (...) {}
                }
            }
        }
        if (inst.mnemonic == "enter")
            return true;
    }
    return false;
}

bool Disassembler::isEpilog(const Instruction& inst) const
{
    return inst.mnemonic == "leave" || inst.mnemonic == "ret" || inst.mnemonic == "retf";
}

bool Disassembler::isReturn(const Instruction& inst) const
{
    return inst.mnemonic == "ret" || inst.mnemonic == "retf" || inst.mnemonic == "iret" || inst.mnemonic == "iretd";
}

bool Disassembler::isTrap(const Instruction& inst) const
{
    return inst.mnemonic == "hlt" || inst.mnemonic == "ud2" || inst.mnemonic == "int3" || inst.mnemonic == "ud0";
}

bool Disassembler::isCall(const Instruction& inst) const
{
    return inst.mnemonic == "call";
}

bool Disassembler::isUnconditionalBranch(const Instruction& inst) const
{
    return inst.mnemonic == "jmp" || inst.mnemonic == "ljmp";
}

bool Disassembler::isConditionalBranch(const Instruction& inst) const
{
    return inst.mnemonic.size() == 2 && inst.mnemonic[0] == 'j' && inst.mnemonic != "jmp";
}

bool Disassembler::isIndirectBranch(const Instruction& inst) const
{
    const string& ops = inst.operands;
    if (inst.mnemonic == "call" || inst.mnemonic == "jmp")
    {
        if (ops.find('[') != string::npos) return true;
        if (ops.size() >= 3 && ops[0] == 'r' && isdigit(ops[1])) return true;
        if (ops.size() >= 4 && ops[0] == 'r' && ops[1] == '1' && isdigit(ops[2])) return true;
        if (ops.size() >= 3 && ops[0] == 'e' && isdigit(ops[1])) return true;
        if (ops[0] == '*') return true;
    }
    return false;
}

addr_t Disassembler::getBranchTarget(const Instruction& inst) const
{
    const string& ops = inst.operands;
    if (ops.empty()) return 0;

    size_t pos = string::npos;
    if (ops[0] == '0' && ops.size() > 2 && ops[1] == 'x')
        pos = 0;
    else if (ops[0] >= '0' && ops[0] <= '9')
        pos = 0;
    else
        pos = ops.find("0x");

    if (pos == string::npos) return 0;

    try {
        return stoull(ops.substr(pos), nullptr, 16);
    } catch (...) {
        return 0;
    }
}

int64_t Disassembler::getRIPDisp(const Instruction& inst) const
{
    if (!m_initialized) return 0;

    auto bytes = inst.bytes;
    if (bytes.empty()) return 0;

    cs_insn* insn = cs_malloc(m_handle);
    if (!insn) return 0;

    const uint8_t* code = bytes.data();
    size_t size = bytes.size();
    addr_t addr = inst.address;
    int64_t disp = 0;

    if (cs_disasm_iter(m_handle, &code, &size, &addr, insn))
    {
        if (insn->detail)
        {
            for (int i = 0; i < insn->detail->x86.op_count; i++)
            {
                const auto& op = insn->detail->x86.operands[i];
                if (op.type == X86_OP_MEM && op.mem.base == X86_REG_RIP)
                {
                    disp = op.mem.disp;
                    break;
                }
            }
        }
    }

    cs_free(insn, 1);
    return disp;
}