module Language.Elna.Architecture.RiscV
    ( BaseOpcode(..)
    , RelocationType(..)
    , Funct3(..)
    , Funct7(..)
    , Funct12(..)
    , Instruction(..)
    , Type(..)
    , XRegister(..)
    , baseOpcode
    , funct3
    , funct12
    , instruction
    , xRegister
    ) where

import qualified Data.ByteString.Builder as ByteString.Builder
import Data.Bits (Bits(..))
import Data.Text (Text)
import Data.Word (Word8, Word32)

data XRegister
    = Zero
    | RA
    | SP
    | GP
    | TP
    | T0
    | T1
    | T2
    | S0
    | S1
    | A0
    | A1
    | A2
    | A3
    | A4
    | A5
    | A6
    | A7
    | S2
    | S3
    | S4
    | S5
    | S6
    | S7
    | S8
    | S9
    | S10
    | S11
    | T3
    | T4
    | T5
    | T6
    deriving Eq

data Funct3
    = ADDI
    | SLTI
    | SLTIU
    | ANDI
    | ORI
    | XORI
    | SLLI
    | SRLI
    | SRAI
    | ADD
    | SLT
    | SLTU
    | AND
    | OR
    | XOR
    | SLL
    | SRL
    | SUB
    | SRA
    | BEQ
    | BNE
    | BLT
    | BLTU
    | BGE
    | BGEU
    | FENCE
    | FENCEI
    | CSRRW
    | CSRRS
    | CSRRC
    | CSRRWI
    | CSRRSI
    | CSRRCI
    | PRIV
    | SB
    | SH
    | SW
    | LB
    | LH
    | LW
    | LBU
    | LHU
    | JALR
    | MUL
    | MULH
    | MULHSU
    | MULHU
    | DIV
    | DIVU
    | REM
    | REMU
    deriving Eq

data Funct12
    = ECALL
    | EBREAK
    deriving Eq

newtype Funct7 = Funct7
    { funct7 :: Word8
    } deriving Eq

data BaseOpcode
    = OpImm
    | Lui
    | Auipc
    | Op
    | Jal
    | Jalr
    | Branch
    | Load
    | Store
    | MiscMem
    | System
    deriving Eq

data Type
    = I XRegister Funct3 XRegister Word32
    | S Word32 Funct3 XRegister XRegister
    | B Word32 Funct3 XRegister XRegister
    | R XRegister Funct3 XRegister XRegister Funct7
    | U XRegister Word32
    | J XRegister Word32
    | Type XRegister Funct3 XRegister Funct12 -- Privileged.
    deriving Eq

data RelocationType
    = RLower12I XRegister Funct3 XRegister Text
    | RLower12S Text Funct3 XRegister XRegister
    | RHigher20 XRegister Text -- Type U.
    | RBranch Text Funct3 XRegister XRegister -- Type B.
    | RJal XRegister Text -- Type J.
    deriving Eq

data Instruction
    = BaseInstruction BaseOpcode Type
    | RelocatableInstruction BaseOpcode RelocationType
    | CallInstruction Text
    deriving Eq

xRegister :: XRegister -> Word8
xRegister Zero = 0
xRegister RA = 1
xRegister SP = 2
xRegister GP = 3
xRegister TP = 4
xRegister T0 = 5
xRegister T1 = 6
xRegister T2 = 7
xRegister S0 = 8
xRegister S1 = 9
xRegister A0 = 10
xRegister A1 = 11
xRegister A2 = 12
xRegister A3 = 13
xRegister A4 = 14
xRegister A5 = 15
xRegister A6 = 16
xRegister A7 = 17
xRegister S2 = 18
xRegister S3 = 19
xRegister S4 = 20
xRegister S5 = 21
xRegister S6 = 22
xRegister S7 = 23
xRegister S8 = 24
xRegister S9 = 25
xRegister S10 = 26
xRegister S11 = 27
xRegister T3 = 28
xRegister T4 = 29
xRegister T5 = 30
xRegister T6 = 31

funct3 :: Funct3 -> Word8
funct3 ADDI = 0b000
funct3 SLTI = 0b001
funct3 SLTIU = 0b011
funct3 ANDI = 0b111
funct3 ORI = 0b110
funct3 XORI = 0b100
funct3 SLLI = 0b000
funct3 SRLI = 0b101
funct3 SRAI = 0b101
funct3 ADD = 0b000
funct3 SLT = 0b010
funct3 SLTU = 0b011
funct3 AND = 0b111
funct3 OR = 0b110
funct3 XOR = 0b100
funct3 SLL = 0b001
funct3 SRL = 0b101
funct3 SUB = 0b000
funct3 SRA = 0b101
funct3 BEQ = 0b000
funct3 BNE = 0b001
funct3 BLT = 0b100
funct3 BLTU = 0b110
funct3 BGE = 0b101
funct3 BGEU = 0b111
funct3 FENCE = 0b000
funct3 FENCEI = 0b001
funct3 CSRRW = 0b001
funct3 CSRRS = 0b010
funct3 CSRRC = 0b011
funct3 CSRRWI = 0b101
funct3 CSRRSI = 0b110
funct3 CSRRCI = 0b111
funct3 PRIV = 0b000
funct3 SB = 0b000
funct3 SH = 0b001
funct3 SW = 0b010
funct3 LB = 0b000
funct3 LH = 0b001
funct3 LW = 0b010
funct3 LBU = 0b100
funct3 LHU = 0b101
funct3 JALR = 0b000
funct3 MUL = 0b000
funct3 MULH = 0b001
funct3 MULHSU = 0b010
funct3 MULHU = 0b011
funct3 DIV = 0b100
funct3 DIVU = 0b101
funct3 REM = 0b110
funct3 REMU = 0b111

funct12 :: Funct12 -> Word8
funct12 ECALL = 0b000000000000
funct12 EBREAK = 0b000000000001

baseOpcode :: BaseOpcode -> Word8
baseOpcode OpImm = 0b0010011
baseOpcode Lui = 0b0110111
baseOpcode Auipc = 0b0010111
baseOpcode Op = 0b0110011
baseOpcode Jal = 0b1101111
baseOpcode Jalr = 0b1100111
baseOpcode Branch = 0b1100011
baseOpcode Load = 0b0000011
baseOpcode Store = 0b0100011
baseOpcode MiscMem = 0b0001111
baseOpcode System = 0b1110011

type' :: Type -> Word32
type' (I rd funct3' rs1 immediate)
    = (fromIntegral (xRegister rd) `shiftL` 7)
    .|. (fromIntegral (funct3 funct3') `shiftL` 12)
    .|. (fromIntegral (xRegister rs1) `shiftL` 15)
    .|. (immediate `shiftL` 20);
type' (S immediate funct3' rs1 rs2)
    = ((immediate .&. 0x1f) `shiftL` 7)
    .|. (fromIntegral (funct3 funct3') `shiftL` 12)
    .|. (fromIntegral (xRegister rs1) `shiftL` 15)
    .|. (fromIntegral (xRegister rs2) `shiftL` 20)
    .|. ((immediate .&. 0xfe0) `shiftL` 20)
type' (B immediate funct3' rs1 rs2)
    = ((immediate .&. 0x800) `shiftR` 4)
    .|. ((immediate .&. 0x1e) `shiftL` 7)
    .|. (fromIntegral (funct3 funct3') `shiftL` 12)
    .|. (fromIntegral (xRegister rs1) `shiftL` 15)
    .|. (fromIntegral (xRegister rs2) `shiftL` 20)
    .|. ((immediate .&. 0x7e0) `shiftL` 20)
    .|. ((immediate .&. 0x1000) `shiftL` 19)
type' (R rd funct3' rs1 rs2 funct7')
    = (fromIntegral (xRegister rd) `shiftL` 7)
    .|. (fromIntegral (funct3 funct3') `shiftL` 12)
    .|. (fromIntegral (xRegister rs1) `shiftL` 15)
    .|. (fromIntegral (xRegister rs2) `shiftL` 20)
    .|. (fromIntegral (funct7 funct7') `shiftL` 25);
type' (U rd immediate)
    = (fromIntegral (xRegister rd) `shiftL` 7)
    .|. (immediate `shiftL` 12)
type' (J rd immediate)
     = (fromIntegral (xRegister rd) `shiftL` 7)
     .|. (immediate .&. 0xff000)
     .|. ((immediate .&. 0x800) `shiftL` 9)
     .|. ((immediate .&. 0x7fe) `shiftL` 20)
     .|. ((immediate .&. 0x100000) `shiftL` 11);
type' (Type rd funct3' rs1 funct12')
    = (fromIntegral (xRegister rd) `shiftL` 7)
    .|. (fromIntegral (funct3 funct3') `shiftL` 12)
    .|. (fromIntegral (xRegister rs1) `shiftL` 15)
    .|. (fromIntegral (funct12 funct12') `shiftL` 20);

relocationType :: RelocationType -> Word32
relocationType (RLower12I rd funct3' rs1 _) = type' $ I rd funct3' rs1 0
relocationType (RLower12S _ funct3' rs1 rs2) = type' $ S 0 funct3' rs1 rs2
relocationType (RHigher20 rd _) = type' $ U rd 0
relocationType (RBranch _ funct3' rs1 rs2) = type' $ B 0 funct3' rs1 rs2
relocationType (RJal rd _) = type' $ J rd 0

instruction :: Instruction -> ByteString.Builder.Builder
instruction = \case
    (BaseInstruction base instructionType) -> go base $ type' instructionType
    (RelocatableInstruction base instructionType) -> go base $ relocationType instructionType
    (CallInstruction _) -> foldMap instruction
        [ BaseInstruction Auipc $ U RA 0
        , BaseInstruction Jalr $ I RA JALR RA 0
        ]
  where
    go base instructionType
        = ByteString.Builder.word32LE
        $ fromIntegral (baseOpcode base)
        .|. instructionType