#ifndef OP_H
#define OP_H

#include "types.h"

/* Total number of registers. */
#define REGISTERS       32
/* Word size of target architecture (RISCV64). */
#define WORD_SIZE       8
/* Tag size for optional/union discriminants. */
#define TAG_SIZE        1
/* Instruction size in bytes (always 32-bit, even on RV64). */
#define INSTR_SIZE      4
/* Stack alignment requirement. */
#define STACK_ALIGNMENT 16
/* The frame pointer register is set as an alias of `S0`. */
#define FP              S0

/* Convenient macro wrappers for `instr`.
 * Some of these, such as BLE and BGT are implemented by swapping the operands
 * of other instructions. */
#define ADDI(rd, rs1, imm)   __instr(I_ADDI, rd, rs1, 0, imm)
#define SLTI(rd, rs1, imm)   __instr(I_SLTI, rd, rs1, 0, imm)
#define SLTIU(rd, rs1, imm)  __instr(I_SLTIU, rd, rs1, 0, imm)
#define XORI(rd, rs1, imm)   __instr(I_XORI, rd, rs1, 0, imm)
#define ORI(rd, rs1, imm)    __instr(I_ORI, rd, rs1, 0, imm)
#define ANDI(rd, rs1, imm)   __instr(I_ANDI, rd, rs1, 0, imm)
#define SLLI(rd, rs1, imm)   __instr(I_SLLI, rd, rs1, 0, imm)
#define SRLI(rd, rs1, imm)   __instr(I_SRLI, rd, rs1, 0, imm)
#define SRAI(rd, rs1, imm)   __instr(I_SRAI, rd, rs1, 0, imm)
#define JALR(rd, rs1, imm)   __instr(I_JALR, rd, rs1, 0, imm)
#define LB(rd, rs1, imm)     __instr(I_LB, rd, rs1, 0, imm)
#define LH(rd, rs1, imm)     __instr(I_LH, rd, rs1, 0, imm)
#define LW(rd, rs1, imm)     __instr(I_LW, rd, rs1, 0, imm)
#define LWU(rd, rs1, imm)    __instr(I_LWU, rd, rs1, 0, imm)
#define LD(rd, rs1, imm)     __instr(I_LD, rd, rs1, 0, imm)
#define LBU(rd, rs1, imm)    __instr(I_LBU, rd, rs1, 0, imm)
#define LHU(rd, rs1, imm)    __instr(I_LHU, rd, rs1, 0, imm)
#define SB(rs2, rs1, imm)    __instr(I_SB, 0, rs1, rs2, imm)
#define SH(rs2, rs1, imm)    __instr(I_SH, 0, rs1, rs2, imm)
#define SW(rs2, rs1, imm)    __instr(I_SW, 0, rs1, rs2, imm)
#define SD(rs2, rs1, imm)    __instr(I_SD, 0, rs1, rs2, imm)
#define BEQ(rs1, rs2, imm)   __instr(I_BEQ, 0, rs1, rs2, imm)
#define BNE(rs1, rs2, imm)   __instr(I_BNE, 0, rs1, rs2, imm)
#define BLT(rs1, rs2, imm)   __instr(I_BLT, 0, rs1, rs2, imm)
#define BGE(rs1, rs2, imm)   __instr(I_BGE, 0, rs1, rs2, imm)
#define BLTU(rs1, rs2, imm)  __instr(I_BLTU, 0, rs1, rs2, imm)
#define BGEU(rs1, rs2, imm)  __instr(I_BGEU, 0, rs1, rs2, imm)
#define BLE(rs1, rs2, imm)   __instr(I_BGE, 0, rs2, rs1, imm)
#define BGT(rs1, rs2, imm)   __instr(I_BLT, 0, rs2, rs1, imm)
#define ADD(rd, rs1, rs2)    __instr(I_ADD, rd, rs1, rs2, 0)
#define SUB(rd, rs1, rs2)    __instr(I_SUB, rd, rs1, rs2, 0)
#define DIV(rd, rs1, rs2)    __instr(I_DIV, rd, rs1, rs2, 0)
#define DIVU(rd, rs1, rs2)   __instr(I_DIVU, rd, rs1, rs2, 0)
#define REM(rd, rs1, rs2)    __instr(I_REM, rd, rs1, rs2, 0)
#define REMU(rd, rs1, rs2)   __instr(I_REMU, rd, rs1, rs2, 0)
#define MUL(rd, rs1, rs2)    __instr(I_MUL, rd, rs1, rs2, 0)
#define SLL(rd, rs1, rs2)    __instr(I_SLL, rd, rs1, rs2, 0)
#define SLT(rd, rs1, rs2)    __instr(I_SLT, rd, rs1, rs2, 0)
#define SLTU(rd, rs1, rs2)   __instr(I_SLTU, rd, rs1, rs2, 0)
#define XOR(rd, rs1, rs2)    __instr(I_XOR, rd, rs1, rs2, 0)
#define SRL(rd, rs1, rs2)    __instr(I_SRL, rd, rs1, rs2, 0)
#define AND(rd, rs1, rs2)    __instr(I_AND, rd, rs1, rs2, 0)
#define OR(rd, rs1, rs2)     __instr(I_OR, rd, rs1, rs2, 0)
#define LUI(rd, imm)         __instr(I_LUI, rd, 0, 0, imm)
#define AUIPC(rd, imm)       __instr(I_AUIPC, rd, 0, 0, imm)
#define JAL(rd, imm)         __instr(I_JAL, rd, 0, 0, imm)
#define JMP(imm)             __instr(I_JMP, 0, 0, 0, imm)
#define MV(rd, rs1)          __instr(I_MV, rd, rs1, 0, 0)
#define NOT(rd, rs1)         __instr(I_NOT, rd, rs1, 0, 0)
#define NEG(rd, rs1)         __instr(I_NEG, rd, rs1, 0, 0)
#define NOP                  __instr(I_NOP, 0, 0, 0, 0)
#define RET                  __instr(I_JALR, ZERO, RA, 0, 0)
#define EBREAK               __instr(I_EBREAK, 0, 0, 0, 0)
#define ECALL                __instr(I_ECALL, 0, 0, 0, 0)
/* RV64I word-width (32-bit) operations */
#define ADDIW(rd, rs1, imm)  __instr(I_ADDIW, rd, rs1, 0, imm)
#define ADDW(rd, rs1, rs2)   __instr(I_ADDW, rd, rs1, rs2, 0)
#define SUBW(rd, rs1, rs2)   __instr(I_SUBW, rd, rs1, rs2, 0)
#define MULW(rd, rs1, rs2)   __instr(I_MULW, rd, rs1, rs2, 0)
#define DIVW(rd, rs1, rs2)   __instr(I_DIVW, rd, rs1, rs2, 0)
#define DIVUW(rd, rs1, rs2)  __instr(I_DIVUW, rd, rs1, rs2, 0)
#define REMW(rd, rs1, rs2)   __instr(I_REMW, rd, rs1, rs2, 0)
#define REMUW(rd, rs1, rs2)  __instr(I_REMUW, rd, rs1, rs2, 0)
#define SLLIW(rd, rs1, imm)  __instr(I_SLLIW, rd, rs1, 0, imm)
#define SRLIW(rd, rs1, imm)  __instr(I_SRLIW, rd, rs1, 0, imm)
#define SRAIW(rd, rs1, imm)  __instr(I_SRAIW, rd, rs1, 0, imm)
#define SLLW(rd, rs1, rs2)   __instr(I_SLLW, rd, rs1, rs2, 0)
#define SRLW(rd, rs1, rs2)   __instr(I_SRLW, rd, rs1, rs2, 0)
#define SRAW(rd, rs1, rs2)   __instr(I_SRAW, rd, rs1, rs2, 0)
/* F Extension - Floating-point instructions */
#define FADD_S(rd, rs1, rs2) __instr(I_FADD_S, rd, rs1, rs2, 0)
#define FSUB_S(rd, rs1, rs2) __instr(I_FSUB_S, rd, rs1, rs2, 0)
#define FMUL_S(rd, rs1, rs2) __instr(I_FMUL_S, rd, rs1, rs2, 0)
#define FDIV_S(rd, rs1, rs2) __instr(I_FDIV_S, rd, rs1, rs2, 0)
#define FEQ_S(rd, rs1, rs2)  __instr(I_FEQ_S, rd, rs1, rs2, 0)
#define FLT_S(rd, rs1, rs2)  __instr(I_FLT_S, rd, rs1, rs2, 0)
#define FLE_S(rd, rs1, rs2)  __instr(I_FLE_S, rd, rs1, rs2, 0)
#define FLW(rd, rs1, imm)    __instr(I_FLW, rd, rs1, 0, imm)
#define FSW(rs2, rs1, imm)   __instr(I_FSW, 0, rs1, rs2, imm)

/* String representations of register names. */
extern const char *reg_names[];

/* Boolean map of caller-saved registers.
 * True for registers that need to be saved by the caller
 * before a function call. */
extern const bool caller_saved_registers[REGISTERS];

/* RISC-V register names. */
typedef enum {
    ZERO = 0,  /* Hard-wired zero */
    RA   = 1,  /* Return address */
    SP   = 2,  /* Stack pointer */
    GP   = 3,  /* Global pointer */
    TP   = 4,  /* Thread pointer */
    T0   = 5,  /* Temporary/alternate link register */
    T1   = 6,  /* Temporary */
    T2   = 7,  /* Temporary */
    S0   = 8,  /* Saved register/frame pointer */
    S1   = 9,  /* Saved register */
    A0   = 10, /* Function arguments/returns */
    A1   = 11,
    A2   = 12, /* Function arguments */
    A3   = 13,
    A4   = 14,
    A5   = 15,
    A6   = 16,
    A7   = 17,
    S2   = 18, /* Saved registers */
    S3   = 19,
    S4   = 20,
    S5   = 21,
    S6   = 22,
    S7   = 23,
    S8   = 24,
    S9   = 25,
    S10  = 26,
    S11  = 27,
    T3   = 28, /* Temporaries */
    T4   = 29,
    T5   = 30,
    T6   = 31
} reg_t;

/* Temporary registers (T1-T6) */
extern const reg_t temp_registers[6];

/* Opcodes for RISC-V base instruction set */
typedef enum {
    OP_LOAD   = 0x03,
    OP_STORE  = 0x23,
    OP_BRANCH = 0x63,
    OP_JALR   = 0x67,
    OP_JAL    = 0x6F,
    OP_OP     = 0x33,
    OP_IMM    = 0x13,
    OP_AUIPC  = 0x17,
    OP_IMM_32 = 0x1B, /* RV64I: ADDIW, SLLIW, SRLIW, SRAIW */
    OP_OP_32 = 0x3B, /* RV64I: ADDW, SUBW, SLLW, SRLW, SRAW, MULW, DIVW, REMW */
    OP_LUI   = 0x37,
    OP_SYSTEM = 0x73,
    OP_FENCE  = 0x0F,
    /* F Extension opcodes */
    OP_LOAD_FP  = 0x07,
    OP_STORE_FP = 0x27,
    OP_OP_FP    = 0x53
} opcode_t;

/* Function3 values */
typedef enum {
    /* Memory operations */
    FUNCT3_BYTE   = 0x0, /* LB/SB - Load/Store Byte */
    FUNCT3_HALF   = 0x1, /* LH/SH - Load/Store Halfword */
    FUNCT3_WORD   = 0x2, /* LW/SW - Load/Store Word */
    FUNCT3_DOUBLE = 0x3, /* LD/SD - Load/Store Doubleword */
    FUNCT3_BYTE_U = 0x4, /* LBU - Load Byte Unsigned */
    FUNCT3_HALF_U = 0x5, /* LHU - Load Halfword Unsigned */
    FUNCT3_WORD_U = 0x6, /* LWU - Load Word Unsigned */

    /* ALU operations */
    FUNCT3_ADD  = 0x0, /* ADD/SUB/ADDI */
    FUNCT3_SLL  = 0x1, /* SLL/SLLI */
    FUNCT3_SLT  = 0x2, /* SLT/SLTI */
    FUNCT3_SLTU = 0x3, /* SLTU/SLTIU */
    FUNCT3_XOR  = 0x4, /* XOR/XORI */
    FUNCT3_SRL  = 0x5, /* SRL/SRA/SRLI/SRAI */
    FUNCT3_OR   = 0x6, /* OR/ORI */
    FUNCT3_AND  = 0x7, /* AND/ANDI */
    /* F Extension function3 codes */
    FUNCT3_WORD_FP = 0x2, /* FLW/FSW - Load/Store Single */
    FUNCT3_FEQ     = 0x2, /* FEQ.S */
    FUNCT3_FLT     = 0x1, /* FLT.S */
    FUNCT3_FLE     = 0x0  /* FLE.S */
} funct3_t;

/* Function7 values */
typedef enum {
    FUNCT7_NORMAL = 0x00,
    FUNCT7_SUB    = 0x20,
    FUNCT7_SRA    = 0x20,
    FUNCT7_MUL    = 0x01,
    /* F Extension function codes */
    FUNCT7_FADD_S = 0x00,
    FUNCT7_FSUB_S = 0x04,
    FUNCT7_FMUL_S = 0x08,
    FUNCT7_FDIV_S = 0x0C,
    FUNCT7_FEQ_S  = 0x50,
    FUNCT7_FLT_S  = 0x50,
    FUNCT7_FLE_S  = 0x50
} funct7_t;

/* Represents a RISC-V instruction in its various formats */
typedef union {
    struct {
        u32 opcode : 7;
        u32 rd     : 5;
        u32 funct3 : 3;
        u32 rs1    : 5;
        u32 rs2    : 5;
        u32 funct7 : 7;
    } r; /* Register format */

    struct {
        u32 opcode   : 7;
        u32 rd       : 5;
        u32 funct3   : 3;
        u32 rs1      : 5;
        u32 imm_11_0 : 12;
    } i; /* Immediate format */

    struct {
        u32 opcode   : 7;
        u32 imm_4_0  : 5;
        u32 funct3   : 3;
        u32 rs1      : 5;
        u32 rs2      : 5;
        u32 imm_11_5 : 7;
    } s; /* Store format */

    struct {
        u32 opcode   : 7;
        u32 imm_11   : 1;
        u32 imm_4_1  : 4;
        u32 funct3   : 3;
        u32 rs1      : 5;
        u32 rs2      : 5;
        u32 imm_10_5 : 6;
        u32 imm_12   : 1;
    } b; /* Branch format */

    struct {
        u32 opcode    : 7;
        u32 rd        : 5;
        u32 imm_31_12 : 20;
    } u; /* Upper immediate format */

    struct {
        u32 opcode    : 7;
        u32 rd        : 5;
        u32 imm_19_12 : 8;
        u32 imm_11    : 1;
        u32 imm_10_1  : 10;
        u32 imm_20    : 1;
    } j; /* Jump format */

    u32 raw; /* Raw 32-bit instruction */
} instr_t;

/* Instruction type. */
typedef enum {
    IFMT_I, /* I-type (immediate) */
    IFMT_R, /* R-type (register) */
    IFMT_S, /* S-type (store) */
    IFMT_B, /* B-type (branch) */
    IFMT_U, /* U-type (upper immediate) */
    IFMT_J, /* J-type (jump) */
} ifmt_t;

/* RISC-V instruction name. */
typedef enum {
    I_LUI,
    I_AUIPC,
    I_JAL,
    I_JALR,
    I_BEQ,
    I_BNE,
    I_BLT,
    I_BGE,
    I_BLTU,
    I_BGEU,
    I_LB,
    I_LH,
    I_LW,
    I_LBU,
    I_LHU,
    I_SB,
    I_SH,
    I_SW,
    I_ADDI,
    I_SLTI,
    I_SLTIU,
    I_XORI,
    I_ORI,
    I_ANDI,
    I_SLLI,
    I_SRLI,
    I_SRAI,
    I_ADD,
    I_SUB,
    I_SLL,
    I_SLT,
    I_SLTU,
    I_XOR,
    I_SRL,
    I_SRA,
    I_OR,
    I_AND,
    I_MUL,
    I_MULH,
    I_MULHSU,
    I_MULHU,
    I_DIV,
    I_DIVU,
    I_REM,
    I_REMU,
    I_MV,
    I_JMP,
    I_NOP,
    I_NOT,
    I_NEG,
    I_EBREAK,
    I_ECALL,
    /* F Extension - Floating-point instructions */
    I_FADD_S,
    I_FSUB_S,
    I_FMUL_S,
    I_FDIV_S,
    I_FEQ_S,
    I_FLT_S,
    I_FLE_S,
    I_FLW,
    I_FSW,
    /* RV64I extensions */
    I_LWU,
    I_LD,
    I_SD,
    I_ADDIW,
    I_SLLIW,
    I_SRLIW,
    I_SRAIW,
    I_ADDW,
    I_SUBW,
    I_SLLW,
    I_SRLW,
    I_SRAW,
    I_MULW,
    I_DIVW,
    I_DIVUW,
    I_REMW,
    I_REMUW
} iname_t;

/* Returns a RISC-V instruction based on the instruction type. */
instr_t instr(iname_t op, reg_t rd, reg_t rs1, reg_t rs2, i32 imm);

static inline instr_t __instr(
    iname_t op, reg_t rd, reg_t rs1, reg_t rs2, i32 imm
) {
    return instr(op, rd, rs1, rs2, imm);
}

/* Return true when a signed 12-bit immediate can encode `value`. */
static inline bool is_small(i32 value) {
    return value >= -2048 && value <= 2047;
}

static inline bool is_branch_imm(i32 value) {
    return value >= -(1 << 12) && value <= ((1 << 12) - 2) && !(value & 1);
}

static inline bool is_jump_imm(i32 value) {
    return value >= -(1 << 20) && value <= ((1 << 20) - 2) && !(value & 1);
}

/* Helper function to sign-extend a value. */
i32 sign_extend(u32 value, int bit_width);
/* Aligns a size to the specified alignment boundary. */
i32 align(i32 size, i32 alignment);
/* Functions to get immediates out of instruction. */
i32 get_i_imm(instr_t instr);
i32 get_s_imm(instr_t instr);
i32 get_b_imm(instr_t instr);
i32 get_j_imm(instr_t instr);

#endif