#include "arm-singlestep.h"

#ifdef ARCH_AARCH64

static void break_imm(unsigned long address, struct thread *th) {
    struct breakpoint *b = add_breakpoint(th->proc, address, 1);
    b->user = 0;
    b->tid = th->id;
    b->enabled = -1;
}

static void break_reg(int reg, struct thread *th) {
    unsigned long address = 0;
    unsigned int *regs = th->state->regs.arm32.regs;

    switch (reg) {
        case ARM_REG_R0:  address = regs[0];  break;
        case ARM_REG_R1:  address = regs[1];  break;
        case ARM_REG_R2:  address = regs[2];  break;
        case ARM_REG_R3:  address = regs[3];  break;
        case ARM_REG_R4:  address = regs[4];  break;
        case ARM_REG_R5:  address = regs[5];  break;
        case ARM_REG_R6:  address = regs[6];  break;
        case ARM_REG_R7:  address = regs[7];  break;
        case ARM_REG_R8:  address = regs[8];  break;
        case ARM_REG_R9:  address = regs[9];  break;
        case ARM_REG_R10: address = regs[10]; break;
        case ARM_REG_R11: address = regs[11]; break;
        case ARM_REG_R12: address = regs[12]; break;
        case ARM_REG_R13: address = regs[13]; break;
        case ARM_REG_R14: address = regs[14]; break;
        case ARM_REG_R15: address = regs[15]; break;
        default: /* todo thread error */ break;
    }

    break_imm(address, th);
}

static int is_pc(csh handle, cs_insn *insn) {
    cs_regs read, write;
    uint8_t read_size, write_size;
    cs_regs_access(handle, insn, read, &read_size, write, &write_size);

    for (uint8_t i = 0; i < write_size; i++) {
        if (write[i] == ARM_REG_PC) {
            return 1;
        }
    }

    return 0;
}

static uint8_t pc_op(cs_arm_op *ops, uint8_t ops_size) {
    uint8_t i;
    for (i = 0; i < ops_size; i++) {
        if (ops[i].type == ARM_OP_REG && ops[i].reg == ARM_REG_PC) {
            break;
        }
    }
    return i;
}

int arm_singlestep(struct thread *th) {
    struct archinfo archinfo;
    struct iovec regs = { &th->state->regs, th->state->regsize };
    architecture_info(&archinfo, &regs);

    int ret = 0;
    csh handle;

    if (cs_open(archinfo.cs_arch, archinfo.cs_mode, &handle) != CS_ERR_OK) {
        /* todo thread error */
        return -1;
    }

    cs_option(handle, CS_OPT_DETAIL, CS_OPT_ON);

    uint64_t address = archinfo.progmctr;
    size_t codez = 128;
    const uint8_t *code = deref(th, address, codez);
    cs_insn *insn = cs_malloc(handle);

    if (cs_disasm_iter(handle, &code, &codez, &address, insn)) {
        if (is_pc(handle, insn)) {
            cs_arm_op *ops = insn->detail->arm.operands;
            uint8_t ops_size = insn->detail->arm.op_count;
            uint8_t pci;

            switch (insn->id) {
                case ARM_INS_B:
                case ARM_INS_BL:
                case ARM_INS_BX:
                case ARM_INS_BLX:
                    if (ops_size == 1) {
                        switch (ops[0].type) {
                            case ARM_OP_REG: break_reg(ops[0].reg, th); break;
                            case ARM_OP_IMM: break_imm(ops[0].imm, th); break;
                            default: ret = -1; /* todo thread error */ break;
                        }
                    } else {
                        ret = -1;
                        /* todo thread error */
                    }
                    break;

                case ARM_INS_POP:
                    if ((pci = pc_op(ops, ops_size)) < ops_size) {
                        unsigned long saddr = archinfo.stackptr + (pci * archinfo.wordsize);
                        unsigned long *sval = deref(th, saddr, archinfo.wordsize);
                        break_imm(*sval, th);
                    } else {
                        ret = -1;
                        /* todo thread error */
                    }
                    break;

                case ARM_INS_MOV:
                    if (pc_op(ops, ops_size) == 0) {
                        if (ops_size == 2) {
                            switch (ops[1].type) {
                                case ARM_OP_REG: break_reg(ops[1].reg, th); break;
                                case ARM_OP_IMM: break_imm(ops[1].imm, th); break;
                                default: ret = -1; /* todo thread error */ break;
                            }
                        } else {
                            ret = -1;
                            /* tr */
                        }
                    } else {
                        ret = -1;
                        /* tr */
                    }
                    break;

                default:
                    ret = -1;
                    /* todo thread error */
                    break;
            }
        }

        /* default case - next sequential instruction */
        break_imm(address, th);
    } else {
        ret = -1;
        /* todo thread error */
    }

    cs_free(insn, 1);
    cs_close(&handle);
    return ret;
}

#endif