//
// Created by bruno on 2.2.2025.
//

#include "core.h"
#include "memory.h"
#include "string.h"

// Initialize CPU
void init_cpu(CPU *cpu) {
    memset(cpu, 0, sizeof(CPU));
    cpu->mode = CPU_MODE_HALTED | CPU_MODE_SECOND;
}

// Helper function for setting flags in the CPU (here we assume bit0 is the Zero flag,
// and bit1 is the Negative flag).
static inline void set_flags(CPU *cpu, int32_t result) {
    cpu->flags = 0;
    if (result == 0)
        cpu->flags |= CPU_FLAG_ZERO;  // Zero flag
    if (result < 0)
        cpu->flags |= CPU_FLAG_NEGATIVE;  // Negative flag
}

// Execute one cycle
void step(CPU *cpu) {
    if (cpu->mode & (CPU_MODE_HALTED | CPU_MODE_PAUSED | CPU_MODE_ERROR)) {
        return;
    }
    if (cpu->pc >= MEM_SIZE || cpu->pc >= cpu->programEnd) {
        if (cpu->mode | CPU_MODE_LOOP) {
            cpu->pc = 0;
        } else {
            cpu->mode |= CPU_MODE_HALTED;
        }
    }
    uint8_t reg1, reg2, imm, temp, temp2;
    uint32_t newPC, addrTemp, oldPC;
    int32_t cmpResult;

    oldPC = cpu->pc;
    newPC = oldPC;

    uint8_t opcode = read_mem(cpu, cpu->pc++);


    const uint32_t differenceAlignment = oldPC % CPU_INSTRUCTION_SIZE;

    if (differenceAlignment) {
        cpu->pc += differenceAlignment;
    }

    switch (opcode) {
        case NOP:
            //Don't do anything
            break;

        case BRK:
            //Pause CPU (for breakpoints)
            cpu->mode |= CPU_MODE_PAUSED;
            break;

        case HLT:
            //Pause CPU (for breakpoints)
            cpu->mode |= CPU_MODE_HALTED;
            break;

        case INC_RN:
            //Increment register
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            write_reg(cpu, reg1, temp + 1);
            break;

        case INC_ADDR:
            //Increment address
            addrTemp = read_address_argument(cpu);
            temp = read_mem(cpu, addrTemp);
            write_mem(cpu, addrTemp, temp + 1);
            break;

        case DEC_RN:
            //Decrement register
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            write_reg(cpu, reg1, temp - 1);
            break;

        case DEC_ADDR:
            //Decrement address
            addrTemp = read_address_argument(cpu);
            temp = read_mem(cpu, addrTemp);
            write_mem(cpu, addrTemp, temp - 1);
            break;

        case MOV_RN_IMM:
            //Load from immediate to register
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            cpu->regs[reg1] = imm;
            break;

        case MOV_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            write_reg(cpu, reg2, temp);
            break;

        case MOV_RN_ADDR:
            // Store register to memory
            reg1 = read_reg_number(cpu);
            addrTemp = read_address_argument(cpu);
            temp = read_reg(cpu, reg1);
            write_mem(cpu, addrTemp, temp);
            break;

        case MOV_ADDR_RN:
            // Load memory to register
            addrTemp = read_address_argument(cpu);
            reg1 = read_reg_number(cpu);
            temp = read_mem(cpu, addrTemp);
            write_reg(cpu, reg1, temp);
            break;

        case SWAP: {
            // Swap contents of two registers.
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp2 = read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp2);
            write_reg(cpu, reg2, temp);
            break;
        }

        case SWAPN: {
            // Swap the nibbles of a register.
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            cpu->regs[reg1] = ((temp & 0x0F) << 4) | ((temp & 0xF0) >> 4);
            write_reg(cpu, reg1, temp);
            break;
        }

        case ADD_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp += read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case ADD_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp += imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case SUB_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp -= read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case SUB_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp -= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case MUL_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp *= read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case MUL_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp *= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case DIV_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp2 = read_reg(cpu, reg2);
            if (temp2 == 0) {
                printf("Error: Division by zero!\n");
                cpu->mode |= CPU_MODE_ERROR;
                return;
            }
            temp /= temp2;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case DIV_RN_IMM:
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            imm = read_mem(cpu, cpu->pc++);
            if (imm == 0) {
                printf("Error: Division by zero!\n");
                cpu->mode |= CPU_MODE_ERROR;
                return;
            }
            temp /= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case MOD_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp2 = read_reg(cpu, reg2);
            if (temp2 == 0) {
                printf("Error: Modulo by zero!\n");
                cpu->mode |= CPU_MODE_ERROR;
                return;
            }
            temp %= temp2;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case MOD_RN_IMM:
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            imm = read_mem(cpu, cpu->pc++);
            if (imm == 0) {
                printf("Error: Modulo by zero!\n");
                cpu->mode |= CPU_MODE_ERROR;
                return;
            }
            temp %= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case NEG_RN:
            reg1 = read_reg_number(cpu);
            temp = -read_reg(cpu, reg1);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case AND_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp &= read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case AND_RN_IMM:
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            imm = read_mem(cpu, cpu->pc++);
            temp &= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case OR_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp |= read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case OR_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp |= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case XOR_RN_RM:
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp ^= read_reg(cpu, reg2);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case XOR_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp ^= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case NOT_RN:
            reg1 = read_reg_number(cpu);
            temp = ~read_reg(cpu, reg1);
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case SHL_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp <<= imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case SHR_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            temp >>= imm;  // Logical right shift (assuming unsigned value)
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case SAR_RN_IMM:
            reg1 = read_reg_number(cpu);
            imm = read_mem(cpu, cpu->pc++);
            temp = read_reg(cpu, reg1);
            // Arithmetic right shift; cast to signed before shifting.
            temp = ((int32_t) temp) >> imm;
            write_reg(cpu, reg1, temp);
            set_flags(cpu, temp);
            break;

        case JMP:
            newPC = read_address_argument(cpu);
            cpu->pc = newPC;
            break;

        case JMP_REL:
            imm = read_mem(cpu, cpu->pc++) & 0xFF;
            cpu->pc += imm;
            break;

        case CMP: {
            // Compare two registers: set flags (Zero, Negative)
            reg1 = read_reg_number(cpu);
            reg2 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            temp2 = read_reg(cpu, reg2);
            cmpResult = (int32_t) temp - (int32_t) temp2;
            set_flags(cpu, cmpResult);
            break;
        }

        case JMP_BIT_SET_RN:
            // Jump if bit in register set
            reg1 = read_reg_number(cpu);
            uint8_t bit = read_mem(cpu, cpu->pc++);
            newPC = read_address_argument(cpu);

            temp = read_reg(cpu, reg1);
            if (reg1 >= REG_COUNT) {
                reg1 = REG_COUNT - 1;
            }
            if (bit > 7) {
                bit = 7;
            }
            if (temp & (1 << bit)) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;

        case JMP_BIT_SET_ADDR: {
            // Jump if bit in register set
            addrTemp = read_address_argument(cpu);
            if (addrTemp >= MEM_SIZE) {
                addrTemp = MEM_SIZE - 1;
            }
            uint8_t bit = read_mem(cpu, cpu->pc++);
            if (bit > 7) {
                bit = 7;
            }
            newPC = read_address_argument(cpu);
            if (read_mem(cpu, addrTemp) & (1 << bit)) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }


        case JE: {
            // Jump if equal (Zero flag set)
            newPC = read_address_argument(cpu);
            if (cpu->flags & CPU_FLAG_ZERO) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JMP_BIT_CLEAR_RN: {
            // Jump if bit in register set
            reg1 = read_reg_number(cpu);
            temp = read_reg(cpu, reg1);
            uint8_t bit = read_mem(cpu, cpu->pc++);
            if (bit > 7) {
                bit = 7;
            }
            newPC = read_address_argument(cpu);
            if (!(temp & (1 << bit))) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JMP_BIT_CLEAR_ADDR: {
            // Jump if bit in register set
            addrTemp = read_address_argument(cpu);
            uint8_t bit = read_mem(cpu, cpu->pc++);
            if (bit > 7) {
                bit = 7;
            }
            newPC = read_address_argument(cpu);
            if (!(read_mem(cpu, addrTemp) & (1 << bit))) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JNE: {
            // Jump if not equal (Zero flag clear)
            newPC = read_address_argument(cpu);
            if (!(cpu->flags & CPU_FLAG_ZERO)) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JG: {
            // Jump if greater: not negative and not zero.
            newPC = read_address_argument(cpu);
            if (!(cpu->flags & CPU_FLAG_NEGATIVE) && !(cpu->flags & CPU_FLAG_ZERO)) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JL: {
            // Jump if less: Negative flag set.
            newPC = read_address_argument(cpu);
            if (cpu->flags & CPU_FLAG_NEGATIVE) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JGE: {
            // Jump if greater or equal: Zero flag set or negative clear.
            newPC = read_address_argument(cpu);
            if ((cpu->flags & CPU_FLAG_ZERO) || !(cpu->flags & CPU_FLAG_NEGATIVE)) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case JLE: {
            // Jump if less or equal: Negative flag set or Zero flag set.
            newPC = read_address_argument(cpu);
            if ((cpu->flags & CPU_FLAG_NEGATIVE) || (cpu->flags & CPU_FLAG_ZERO)) {
                cpu->pc = newPC;
                break;
            }
            cpu->pc += CPU_INSTRUCTION_SIZE;
            break;
        }

        case CALL: {
            newPC = read_address_argument(cpu);
            write_stack(cpu);
            cpu->pc = newPC;
            break;
        }

        case RET:
            // For RET, assume that the return address was stored on the stack.
            read_stack(cpu);
            break;

        default:
            printf("Unknown opcode: %d\n", opcode);
            cpu->mode |= CPU_MODE_ERROR;
    }
    if (oldPC == newPC) {
        const uint32_t remainingBytes = CPU_INSTRUCTION_SIZE - (cpu->pc - oldPC);
        if (remainingBytes > CPU_INSTRUCTION_SIZE) {
            printf("HELP, INSTRUCTION SIZE SMALLER THAN INSTRUCTION");
        }
        cpu->pc += remainingBytes;
    }
    cpu->cycle++;
}