#include "InstHandler.h"

#define HANDLER(f) [this](inst_t* inst){return f(inst);}

#define REQ_TKS(n) if (inst->tokens.size() != n) return 0
#define PARSE_REG(n, r) if (!parseRegister(inst->tokens[n], r)) return 0
#define PARSE_FPREG(n, r) if (!parseFPRegister(inst->tokens[n], r)) return 0
#define PARSE_IMM(n, i, d) if (!parseImmediate(inst->tokens[n], i, d)) return 0

#define OPCODE(o) ((uint32_t)o << 26)
#define RS(r) (r << 21)
#define RT(r) (r << 16)
#define RD(r) (r << 11)
#define IMM16(i) (i & 0x0000FFFF)

#define COMB_I(o) 0u | OPCODE(o) | RS(rs) | RT(rt) | IMM16(imm)
#define COMB_R(o, f) 0u | OPCODE(o) | RS(rs) | RT(rt) | RD(rd) | ((uint32_t)f & 0x0000003F)


// set up function map
InstHandler::InstHandler(std::vector<label_t>* labels_) :
  labels(labels_)
{
    funcMap["ld"] = HANDLER(handleLD);
    funcMap["l.d"] = HANDLER(handleLfD);
    funcMap["sd"] = HANDLER(handleSD);
    funcMap["s.d"] = HANDLER(handleSfD);
    funcMap["daddi"] = HANDLER(handleDADDI);
    funcMap["daddiu"] = HANDLER(handleDADDIU);
    funcMap["beq"] = HANDLER(handleBEQ);
    funcMap["bne"] = HANDLER(handleBNE);
    funcMap["dadd"] = HANDLER(handleDADD);
    funcMap["dsub"] = HANDLER(handleDSUB);
    funcMap["add.d"] = HANDLER(handleADDfD);
    funcMap["sub.d"] = HANDLER(handleSUBfD);
    funcMap["mul.d"] = HANDLER(handleMULfD);
    funcMap["div.d"] = HANDLER(handleDIVfD);
    funcMap["j"] = HANDLER(handleJ);
    funcMap["halt"] = HANDLER(handleHALT);
    funcMap["nop"] = HANDLER(handleNOP);
    funcMap["dump"] = HANDLER(handleDUMP);
}

// call the appropriate function. this is where the instruction op is validated
int InstHandler::handle(inst_t* inst) {
    try {
        return funcMap.at(inst->tokens[0])(inst);
    } catch (...) {
        // not found in map
        return 0;
    }
}


int InstHandler::handleLD(inst_t* inst) {
    REQ_TKS(4);
    
    uint32_t rt, rs;
    PARSE_REG(1, rt);
    PARSE_REG(3, rs);

    uint32_t imm;
    PARSE_IMM(2, imm, false);

    inst->hex = COMB_I(55);
    return 1;
}

int InstHandler::handleLfD(inst_t* inst) {
    REQ_TKS(4);
    
    uint32_t rt, rs;
    PARSE_FPREG(1, rt);
    PARSE_REG(3, rs);

    uint32_t imm;
    PARSE_IMM(2, imm, false);

    inst->hex = COMB_I(53);
    return 1;
}

int InstHandler::handleSD(inst_t* inst) {
    REQ_TKS(4);
    
    uint32_t rt, rs;
    PARSE_REG(1, rt);
    PARSE_REG(3, rs);

    uint32_t imm;
    PARSE_IMM(2, imm, false);

    inst->hex = COMB_I(63);
    return 1;
}

int InstHandler::handleSfD(inst_t* inst) {
    REQ_TKS(4);
    
    uint32_t rt, rs;
    PARSE_FPREG(1, rt);
    PARSE_REG(3, rs);

    uint32_t imm;
    PARSE_IMM(2, imm, false);

    inst->hex = COMB_I(61);
    return 1;
}

int InstHandler::handleDADDI(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rt, rs;
    PARSE_REG(1, rt);
    PARSE_REG(2, rs);

    uint32_t imm;
    PARSE_IMM(3, imm, false);

    inst->hex = COMB_I(24);
    return 1;
}

int InstHandler::handleDADDIU(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rt, rs;
    PARSE_REG(1, rt);
    PARSE_REG(2, rs);

    uint32_t imm;
    PARSE_IMM(3, imm, false);

    inst->hex = COMB_I(25);
    return 1;
}

int InstHandler::handleBEQ(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rt, rs;
    PARSE_REG(1, rt);
    PARSE_REG(2, rs);

    uint32_t imm;
    if (!parseImmediateBranch(inst->tokens[3], imm, inst->loc)) return 0;

    inst->hex = COMB_I(4);
    return 1;
}

int InstHandler::handleBNE(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rt, rs;
    PARSE_REG(1, rt);
    PARSE_REG(2, rs);

    uint32_t imm;
    if (!parseImmediateBranch(inst->tokens[3], imm, inst->loc)) return 0;

    inst->hex = COMB_I(5);
    return 1;
}

int InstHandler::handleDADD(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rd, rs, rt;
    PARSE_REG(1, rd);
    PARSE_REG(2, rs);
    PARSE_REG(3, rt);

    inst->hex = COMB_R(0, 44);
    return 1;
}

int InstHandler::handleDSUB(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rd, rs, rt;
    PARSE_REG(1, rd);
    PARSE_REG(2, rs);
    PARSE_REG(3, rt);

    inst->hex = COMB_R(0, 46);
    return 1;
}

int InstHandler::handleADDfD(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rd, rs, rt;
    PARSE_FPREG(1, rd);
    PARSE_FPREG(2, rs);
    PARSE_FPREG(3, rt);

    inst->hex = COMB_R(0, 47);
    return 1;
}

int InstHandler::handleSUBfD(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rd, rs, rt;
    PARSE_FPREG(1, rd);
    PARSE_FPREG(2, rs);
    PARSE_FPREG(3, rt);

    inst->hex = COMB_R(0, 48);
    return 1;
}

int InstHandler::handleMULfD(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rd, rs, rt;
    PARSE_FPREG(1, rd);
    PARSE_FPREG(2, rs);
    PARSE_FPREG(3, rt);

    inst->hex = COMB_R(0, 49);
    return 1;
}

int InstHandler::handleDIVfD(inst_t* inst) {
    REQ_TKS(4);

    uint32_t rd, rs, rt;
    PARSE_FPREG(1, rd);
    PARSE_FPREG(2, rs);
    PARSE_FPREG(3, rt);

    inst->hex = COMB_R(0, 50);
    return 1;
}

int InstHandler::handleJ(inst_t* inst) {
    REQ_TKS(2);

    uint32_t imm;
    if (!parseImmediateJump(inst->tokens[1], imm)) return 0;

    inst->hex = 0u | OPCODE(2) | (imm & 0x03FFFFFF);
    return 1;
}

int InstHandler::handleHALT(inst_t* inst) {
    REQ_TKS(1);

    inst->hex = 0u | OPCODE(1);
    return 1;
}

int InstHandler::handleNOP(inst_t* inst) {
    REQ_TKS(1);

    inst->hex = 0u | OPCODE(3);
    return 1;
}

int InstHandler::handleDUMP(inst_t* inst) {
    REQ_TKS(2);

    uint32_t imm;
    PARSE_IMM(1, imm, true);

    inst->hex = 0u | OPCODE(44) | (imm & 0x03FFFFFF);
    return 1;
}

// verifies register arguments are written correctly and sets the reg var to the reg number
int InstHandler::parseRegister(std::string arg, uint32_t& reg) {
    static std::regex regRegex("^[rR](\\d{1,2})$");
    if (!std::regex_match(arg, matches, regRegex)) return 0;
    reg = std::stoul(matches.str(1));
    if (reg > 31) return 0;
    return 1;
}

int InstHandler::parseFPRegister(std::string arg, uint32_t& reg) {
    static std::regex regRegex("^[fF](\\d{1,2})$");
    if (!std::regex_match(arg, matches, regRegex)) return 0;
    reg = std::stoul(matches.str(1));
    if (reg > 31) return 0;
    return 1;
}

// for most immediate functions; parses the immediate value or looks up the label value
int InstHandler::parseImmediate(std::string arg, uint32_t& imm, bool isDump) {
    static std::regex immRegex("^-?\\d+$");
    static std::regex labelRegex("^[[:alnum:]]{1,12}$");
    
    if (std::regex_match(arg, immRegex)) {
        try {
            imm = std::stol(arg);
            return 1;
        } catch (...) {
            return 0;
        }
    } else if (!isDump && std::regex_match(arg, labelRegex)) {
        imm = findLabel(arg);
        if (imm == (uint32_t)-1) return 0;
        imm <<= 2;
        return 1;
    } else {
        return 0;
    }
}

// jumps don't shift labels left at all, but also aren't pc-relative?
int InstHandler::parseImmediateJump(std::string arg, uint32_t& imm) {
    static std::regex immRegex("^-?\\d+$");
    static std::regex labelRegex("^[[:alnum:]]{1,12}$");

    if (std::regex_match(arg, immRegex)) {
        try {
            imm = std::stol(arg);
            return 1;
        } catch (...) {
            return 0;
        }
    } else if (std::regex_match(arg, labelRegex)) {
        imm = findLabel(arg);
        if (imm == (uint32_t)-1) return 0;
        return 1;
    } else {
        return 0;
    }
}

// branches to labels are pc-relative i guess
int InstHandler::parseImmediateBranch(std::string arg, uint32_t& imm, uint32_t icount) {
    static std::regex immRegex("^-?\\d+$");
    static std::regex labelRegex("^[[:alnum:]]{1,12}$");

    if (std::regex_match(arg, immRegex)) {
        try {
            imm = std::stol(arg);
            return 1;
        } catch (...) {
            return 0;
        }
    } else if (std::regex_match(arg, labelRegex)) {
        imm = findLabel(arg);
        if (imm == (uint32_t)-1) return 0;
        imm -= icount + 1;
        return 1;
    } else {
        return 0;
    }
}

// looks up label
uint32_t InstHandler::findLabel(std::string l) {
    for (uint64_t i = 0; i < labels->size(); i++) {
        if (labels->at(i).name == l) return labels->at(i).loc;
    }
    return -1;
}