from .instructions import Instruction, InstructionSet
from ctypes import c_uint32

RV32I = InstructionSet()

class R(Instruction):
    funct3  = None
    funct7  = None
    opcode  = None
    def __init__(self, rd, rs1, rs2):
        self.rd  = rd
        self.rs1 = rs1
        self.rs2 = rs2

    def compile(self):
        # TODO: ensure sizes and convert register names to number...
        return c_uint32(
            (self.funct7  << 25) +\
            (self.rs2     << 20) +\
            (self.rs1     << 15) +\
            (self.funct3  << 12) +\
            (self.rd      <<  7) +\
            self.opcode
        )

class I(Instruction):
    funct3 = None
    opcode = None

    def __init__(self, rd, rs, imm):
        self.rd = rd
        self.rs = rs
        self.imm = imm

    def compile(self):
        return c_uint32(
            (self.imm    << 20) +\
            (self.rs     << 15) +\
            (self.funct3 << 12) +\
            (self.rd     <<  7) +\
            self.opcode
        )

class S(Instruction):
    pass

class B(Instruction):
    pass

class U(Instruction):
    pass

class J(Instruction):
    # TODO check if it's correct
    opcode = None

    def __init__(self, rd, imm):
        self.rd = rd
        self.imm = imm

    def compile(self):
        return c_uint32(
            (self.imm    << 12) +\
            (self.rd     <<  7) +\
            self.opcode
        )


@RV32I.instruction
class add(R):
    name    = "add"
    opcode  = 0b0110011
    funct3  = 0b000
    funct7  = 0b0000000

    def execute(self):
        return pc + self.size

@RV32I.instruction
class addi(I):
    name    = "addi"
    opcode  = 0b0010011
    funct3  = 0b000

@RV32I.instruction
class jal(J):
    name    = "jal"
    opcode  = 0b1101111

    def execute(self):
        # TODO
        # - Save current pc in rd
        # - Make pc from `imm`
        # - Return new pc
        return pc


@RV32I.pseudoinstruction
class j(J):
    name    = "j"
    def __new__(cls, imm):
        return jal("x0", imm)


if __name__ == "__main__":
    print(RV32I)
    print(RV32I.instructions)