#include <stdlib.h>
#include <assert.h>
#include "execute.h"
#include "print.h"

// Defines the maximum value that can be held in a register
#define MAX_REG_VAL ((1 << DWORD_BITS) - 1)

// The number of bits to shift the immediate value in an arithmetic immediate data processing
// instruction if the shift flag is enabled.
#define DPI_ARITHM_SHIFT 12 

// The number of bits to shift the immediate value in a wide move immediate data processing
// instruction if the shift flag is enabled.
#define DPI_WIDEMOV_SHIFT 16

// Prototypes
void execute_SDT(Machine *state, a64inst_instruction *inst);
void execute_Branch(Machine *state, a64inst_instruction *inst);
void executeMultiply(Machine *state, a64inst_instruction *inst);

// Return maximum of two dwords
static dword max(dword a, dword b) {
    return a > b ? a : b;
}

// Truncate a given value to the size of a word or dword depending on the register type
static dword truncateValue(dword value, a64inst_regType regType) {
    if (regType == a64inst_X) {
        return value;
    } else {
        return (word)value;
        //return value & (dword)(((dword)1 << WORD_BITS) - 1);
    }
}

// Sign extend a given value to a 64-bit signed integer given the number of bits
static int64_t signExtend(dword value, unsigned int n) {
    if (n == 0 || n >= 64) {
        // If n_bits is 0 or greater than or equal to 64, return the value as is
        return (int64_t)value;
    }
    
    uint64_t sign_bit_mask = (uint64_t)1 << (n - 1);
    
    // Mask to isolate the n-bit value
    uint64_t n_bit_mask = (sign_bit_mask << 1) - 1;

    // Check if the sign bit is set
    if (value & sign_bit_mask) {
        // Sign bit is set, extend the sign
        return (int64_t)(value | ~n_bit_mask);
    } else {
        // Sign bit is not set, return the value as is
        return (int64_t)(value & n_bit_mask);
    }
}

// Read from processor register, ensuring that a valid register specifier is given
// and accounting for the case where the zero register is accessed. Truncate
// the 32 most significant bits stored in the R register when reading W register.
static dword readRegister(Machine *state, a64inst_regSpecifier reg, a64inst_regType regType) {
    assert(reg <= REGISTER_COUNT);
    if (reg == ZERO_REGISTER) {
        return 0;
    } else {
        return truncateValue(state->registers[reg], regType);
    }
}

// TODO:

// Write to a processor register, ensuring that a valid register specifier is given
// and truncating the value being written when it can't fit in the specified register
static void writeRegister(Machine *state, a64inst_regSpecifier reg, a64inst_regType regType, dword value) {
    assert(reg <= REGISTER_COUNT);
    if (reg != ZERO_REGISTER) {
        state->registers[reg] = truncateValue(value, regType);
    }
}

// Returns the position of the MSB of the given register type
inline static dword getMSBPos(a64inst_regType regType) {
    return (regType ? DWORD_BITS : WORD_BITS) - 1;
}

// Returns the MSB of the given value assuming it's of the size stored in the given register type
inline static uint8_t getMSB(dword value, a64inst_regType regType) {
    return value >> getMSBPos(regType); 
}

// Updates N and Z condition codes given the machine and a result value
static void updateCondNZ(Machine *state, dword result, a64inst_regType regType) {
    state->conditionCodes.Negative = getMSB(result, regType);
    state->conditionCodes.Zero = result == 0;
}

// Execute a data processing immediate instruction
static void executeDPImmediate(Machine *state, a64inst_instruction *inst) {
    assert(inst->type == a64inst_DPIMMEDIATE);
    
    a64inst_regType regType = inst->data.DPImmediateData.regType;
    a64inst_regSpecifier dest = inst->data.DPImmediateData.dest;
    switch(inst->data.DPImmediateData.DPIOpType) {
        
        // Execute an arithmetic immediate data processing instruction
        case a64inst_DPI_ARITHM:;

            // If shift flag is enabled, logical left shift by the number of bits specified by the architecture
            dword arithmImm = inst->data.DPImmediateData.processOpData.arithmData.immediate;
            dword srcVal = state->registers[inst->data.DPImmediateData.processOpData.arithmData.src];
            if (inst->data.DPImmediateData.processOpData.arithmData.shiftImmediate) {
                arithmImm = truncateValue(arithmImm << DPI_ARITHM_SHIFT, regType); 
            }

            switch(inst->data.DPImmediateData.processOp) {
                
                dword result;
                case(a64inst_ADDS):
                    result = srcVal + arithmImm;
                    writeRegister(state, dest, regType, result);
                    
                    updateCondNZ(state, result, regType);
                    state->conditionCodes.Overflow = max(srcVal, arithmImm) > result;
                    state->conditionCodes.Carry = state->conditionCodes.Overflow;
                    break;
                
                case(a64inst_ADD):
                    writeRegister(state, dest, regType, srcVal + arithmImm);
                    break;
                
                case(a64inst_SUBS):
                    result = srcVal - arithmImm;
                    writeRegister(state, dest, regType, result);
                    
                    updateCondNZ(state, result, regType);
                    state->conditionCodes.Overflow = srcVal < result;
                    state->conditionCodes.Carry = state->conditionCodes.Overflow;
                    break;
                
                case(a64inst_SUB):
                    writeRegister(state, dest, regType, srcVal - arithmImm);
                    break;

                // Unknown opcode detected!
                default:
                    fprintf(stderr, "Unknown opcode detected in a DPI arithmetic instruction!\n");
                    break;
            }
            break;

        // Execute a wide move immediate data processing instruction
        case a64inst_DPI_WIDEMOV:;
            uint8_t shiftScalar = inst->data.DPImmediateData.processOpData.wideMovData.shiftScalar; 
            dword wideMovImm = inst->data.DPImmediateData.processOpData.wideMovData.immediate;
            
            // NOTE: Not checking that shiftScalar has valid value for 32bit registers. Possibly add explicit error.
            //printf("%x\n", wideMovImm << (shiftScalar * DPI_WIDEMOV_SHIFT) & );
            wideMovImm = truncateValue(wideMovImm << (shiftScalar * DPI_WIDEMOV_SHIFT), regType);
            switch(inst->data.DPImmediateData.processOp) {

                case(a64inst_MOVN):
                    writeRegister(state, dest, regType, ~wideMovImm);
                    break;

                case(a64inst_MOVZ):
                    writeRegister(state, dest, regType, wideMovImm);
                    break;

                case(a64inst_MOVK):;
                    dword result = readRegister(state, dest, regType);
                    result = (result & ~(((1 << DPI_WIDEMOV_SHIFT) - 1) << shiftScalar)) | wideMovImm;
                    writeRegister(state, dest, regType, result); 
                    break;
                
                default:
                    fprintf(stderr, "Unknown opcode detected in a DPI wide move instruction!\n");
                    break;
            }
            break;

        // Unknown instruction detected!    
        default:
            fprintf(stderr, "Attempting to execute instruction with unknown DPI operand type!\n");
            break;
    }
}

// Execute a data processing register instruction
static void executeDPRegister(Machine *state, a64inst_instruction *inst) {
    assert(inst->type == a64inst_DPREGISTER);

    a64inst_regType regType = inst->data.DPRegisterData.regType;
    a64inst_regSpecifier dest = inst->data.DPRegisterData.dest;
    dword src1Val = readRegister(state, inst->data.DPRegisterData.src1, regType);
    dword src2Val = readRegister(state, inst->data.DPRegisterData.src2, regType);

    switch(inst->data.DPRegisterData.DPROpType) {

        // Execute an arithmetic or logic register data processing instruction
        case a64inst_DPR_ARITHMLOGIC:;

            // Apply shift to value held in second register
            a64inst_DPRegister_ArithmLogicData *arithmLogicData = &inst->data.DPRegisterData.processOpData.arithmLogicData;
            uint8_t shiftAmount = arithmLogicData->shiftAmount;
            switch(arithmLogicData->shiftType) {

                case a64inst_LSL:
                    src2Val = truncateValue(src2Val << shiftAmount, regType);
                    break;
  
                case a64inst_LSR:
                    src2Val = truncateValue(src2Val >> shiftAmount, regType);
                    break;
  
                case a64inst_ASR:
                    if (regType == a64inst_X) {
                        src2Val = truncateValue((int64_t)src2Val >> shiftAmount, regType);
                    } else {
                        src2Val = truncateValue((int32_t)src2Val >> shiftAmount, regType);
                    }
                    break;
  
                case a64inst_ROR:
                    if (arithmLogicData->type != a64inst_DPR_LOGIC) {
                        fprintf(stderr, "Attempting to perform ROR shift on non-logic register data processing instruction!\n");
                    }
                    src2Val = truncateValue(src2Val >> shiftAmount | src2Val << (getMSBPos(regType) - shiftAmount), regType);
                    break;

                default:
                    fprintf(stderr, "Attempting to execute arithmetic/logic register data processing instruction with invalid shift type!\n");
                    break;
            }

            // Negate second operand if negShiftedSrc2 flag is enabled
            if (arithmLogicData->negShiftedSrc2) {
                src2Val = truncateValue(~src2Val, regType);
            }
            
            dword result;
            switch(arithmLogicData->type) {

                case a64inst_DPR_ARITHM:
                    switch(inst->data.DPRegisterData.processOp) {
                    
                        case(a64inst_ADDS):
                            result = src1Val + src2Val;
                            writeRegister(state, dest, regType, result);
                            
                            updateCondNZ(state, result, regType);
                            state->conditionCodes.Overflow = max(src1Val, src2Val) > result;
                            state->conditionCodes.Carry = state->conditionCodes.Overflow;
                            break;
                        
                        case(a64inst_ADD):
                            writeRegister(state, dest, regType, src1Val + src2Val);
                            break;
                        
                        case(a64inst_SUBS):
                            result = src1Val - src2Val;
                            writeRegister(state, dest, regType, result);
                            
                            updateCondNZ(state, result, regType);
                            state->conditionCodes.Overflow = src1Val < result;
                            state->conditionCodes.Carry = state->conditionCodes.Overflow;
                            break;
                        
                        case(a64inst_SUB):
                            writeRegister(state, dest, regType, src1Val - src2Val);
                            break;

                        // Unknown opcode detected!
                        default:
                            fprintf(stderr, "Unknown opcode detected in a DPI arithmetic instruction!\n");
                            break;
                    }
                    break;

                case a64inst_DPR_LOGIC:
                    switch(inst->data.DPRegisterData.processOp) {
                        
                        case a64inst_AND:
                            writeRegister(state, dest, regType, src1Val & src2Val);
                            break;

                        case a64inst_OR:
                            writeRegister(state, dest, regType, src1Val | src2Val);
                            break;

                        case a64inst_XOR:
                            writeRegister(state, dest, regType, src1Val ^ src2Val);
                            break;
                        
                        case a64inst_AND_FLAGGED:;
                            result = src1Val & src2Val;
                            writeRegister(state, dest, regType, result);
                            state->conditionCodes.Overflow = 0;
                            state->conditionCodes.Carry = 0;
                            updateCondNZ(state, result, regType);
                            break;
                    }
                    break;

                default:
                    fprintf(stderr, "Attempting to execute an instruction with an unknown DPR arithmetic or logic subtype!\n");
                    break;                
            }
            break;

        // Execute a multiply register data processing instruction
        case a64inst_DPR_MULTIPLY:
            break;

        // Unknown instruction detected!    
        default:
            fprintf(stderr, "Attempting to execute instruction with unknown DPR operand type!\n");
            break;
    }
}

void execute(Machine *state, a64inst_instruction *inst) {

    switch (inst->type) {

        // Halt the program
        case a64inst_HALT:
            break;

        // Execute a data processing immediate instruction
        case a64inst_DPIMMEDIATE:
            executeDPImmediate(state, inst);
            break;

        // Execute a branch instruction
        case a64inst_BRANCH:
            execute_Branch(state, inst);
            break;

        // Execute a data processing register instruction
        case a64inst_DPREGISTER:
            if (inst->data.DPRegisterData.DPROpType == a64inst_DPR_MULTIPLY)
                executeMultiply(state, inst);
            else
                executeDPRegister(state, inst);
            break;

        case a64inst_SINGLETRANSFER: 
            execute_SDT(state, inst);
            break;

        // Unknown instruction
        default:
            break;
    }

}

void execute_SDT(Machine *state, a64inst_instruction *inst) {
    word address;
    bool isLoad;
    if (inst->data.SingleTransferData.SingleTransferOpType == a64inst_SINGLE_TRANSFER_LOAD_LITERAL) {
        // Load Literal
        isLoad = true;
        address = state->pc + inst->data.SingleTransferData.processOpData.loadLiteralData.offset * 4;
    } else {
        address = state->registers[inst->data.SingleTransferData.processOpData.singleDataTransferData.base];
        isLoad = inst->data.SingleTransferData.processOpData.singleDataTransferData.transferType == a64inst_LOAD;
        switch (inst->data.SingleTransferData.processOpData.singleDataTransferData.addressingMode) {
            case a64inst_UNSIGNED_OFFSET:
                address += inst->data.SingleTransferData.processOpData.singleDataTransferData.a64inst_addressingModeData.unsignedOffset * (inst->data.SingleTransferData.regType == a64inst_W ? 4 : 8);
                break;
            case a64inst_REGISTER_OFFSET:
                address += state->registers[inst->data.SingleTransferData.processOpData.singleDataTransferData.a64inst_addressingModeData.offsetReg];
                break;
            case a64inst_PRE_INDEXED:
                address += inst->data.SingleTransferData.processOpData.singleDataTransferData.a64inst_addressingModeData.indexedOffset;
                state->registers[inst->data.SingleTransferData.processOpData.singleDataTransferData.base] = address;
                break;
            case a64inst_POST_INDEXED:
                state->registers[inst->data.SingleTransferData.processOpData.singleDataTransferData.base] = address + inst->data.SingleTransferData.processOpData.singleDataTransferData.a64inst_addressingModeData.indexedOffset;
                break;
        }
    }

    if (isLoad) {
        if (inst->data.SingleTransferData.regType == a64inst_W) {
            // 32 bit access
            state->registers[inst->data.SingleTransferData.target] = readWord(state->memory, address);
        } else {
            state->registers[inst->data.SingleTransferData.target] = readDoubleWord(state->memory, address);
        }
    } else {
        if (inst->data.SingleTransferData.regType == a64inst_W) {
            // 32 bit access
            state->registers[inst->data.SingleTransferData.target] = readWord(state->memory, address);
        } else {
            state->registers[inst->data.SingleTransferData.target] = readDoubleWord(state->memory, address);
        }
    }

}

static bool isConditionMet(Machine* state, a64inst_ConditionType cond) {
    switch(cond) {
        case EQ:
            return state->conditionCodes.Zero;
        case NE:
            return !state->conditionCodes.Zero;
        case GE:
            return state->conditionCodes.Negative == state->conditionCodes.Overflow;
        case LT:
            return state->conditionCodes.Negative != state->conditionCodes.Overflow;
        case GT:
            return !state->conditionCodes.Zero && (state->conditionCodes.Negative == state->conditionCodes.Overflow);
        case LE:
            return state->conditionCodes.Zero || (state->conditionCodes.Negative != state->conditionCodes.Overflow);
        case AL:
            return true;
        default:
            fprintf(stderr, "Unknown condition specified!\n");
            exit(1);
    }
}

void execute_Branch(Machine *state, a64inst_instruction *inst) {
    switch (inst->data.BranchData.BranchType) {
        case a64inst_UNCONDITIONAL:
            state->pc += signExtend(inst->data.BranchData.processOpData.unconditionalData.unconditionalOffset * 4, 26);
            break;
            
        case a64inst_REGISTER:
            state->pc = state->registers[inst->data.BranchData.processOpData.registerData.src];
            break;

        case a64inst_CONDITIONAL:
            if (isConditionMet(state, inst->data.BranchData.processOpData.conditionalData.cond)) {
                state->pc += signExtend(inst->data.BranchData.processOpData.conditionalData.offset * 4, 19);
            }
            break;
    }
}

void executeMultiply(Machine *state, a64inst_instruction *inst) {
    dword product = state->registers[inst->data.DPRegisterData.src1] * state->registers[inst->data.DPRegisterData.src2];
    dword sum = readRegister(state, inst->data.DPRegisterData.processOpData.multiplydata.summand, inst->data.DPRegisterData.regType) + (inst->data.DPRegisterData.processOpData.multiplydata.negProd ? -product : product);
    writeRegister(state, inst->data.DPRegisterData.dest, inst->data.DPRegisterData.regType, sum);
}