{-# LANGUAGE DataKinds #-}
{-# LANGUAGE NumericUnderscores #-}

module Execute(execute) where

import Clash.Prelude
import Decode(DecodeResult(..))
import DecodeTypes(
    Opcode(..),
    RTypeFields(..), ITypeFields(..), STypeFields(..),
    BTypeFields(..), UTypeFields(..), JTypeFields(..)
  )
import Types(Addr, DoubleWord)
import Exceptions(Exception(..))
import BusTypes(
    WriteRequest(..),
    ReadRequest(..),
    TransactionSize(..),
    BusVal(..)
  )
import RegFiles(RegFileIdx, RegVal(..))
import Util((|>))

data ExecuteResult  = ReadRequest {readRequest :: ReadRequest, insnAddr :: Addr}
                    | WriteRequest {writeRequest :: WriteRequest, insnAddr :: Addr}
                    | WriteBackGPR  {idx :: RegFileIdx, val :: DoubleWord}
                    | Jump         { targetAddr :: Addr }
                    | DecodeException {exception :: Exception, insnAddr :: Addr}
                    | InstructionException {exception :: Exception, insnAddr :: Addr}
                    deriving (Generic, Show, Eq, NFDataX)

-- Helper functions to extract values from RegVal
extractRegVal :: RegVal -> DoubleWord
extractRegVal (Value _ val) = val
extractRegVal (Unpopulated _) = undefined

-- Execute function
execute :: DecodeResult -> ExecuteResult
execute (Opcode opcode addr) = case opcode of
  -- R-Type Instructions
  ADD  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 + val2)
  SUB  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 - val2)
  XOR  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 `xor` val2)
  OR   (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 .|. val2)
  AND  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 .&. val2)
  SLL  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 `shiftL` fromIntegral (val2 .&. 0x3F))
  SRL  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (val1 `shiftR` fromIntegral (val2 .&. 0x3F))
  SRA  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = unpack (pack (extractRegVal rs1) :: BitVector 64) :: Signed 64
        val2 = extractRegVal rs2
    in WriteBackGPR rd (bitCoerce (val1 `shiftR` fromIntegral (val2 .&. 0x3F)))
  SLT  (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = unpack (pack (extractRegVal rs1) :: BitVector 64) :: Signed 64
        val2 = unpack (pack (extractRegVal rs2) :: BitVector 64) :: Signed 64
    in WriteBackGPR rd (if val1 < val2 then 1 else 0)
  SLTU (RTypeFields rd _ rs1 rs2 _) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
    in WriteBackGPR rd (if val1 < val2 then 1 else 0)

  -- I-Type Instructions
  ADDI  (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        immVal = signExtend imm
    in WriteBackGPR rd (val1 + immVal)
  XORI  (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        immVal = signExtend imm
    in WriteBackGPR rd (val1 `xor` immVal)
  ORI   (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        immVal = signExtend imm
    in WriteBackGPR rd (val1 .|. immVal)
  ANDI  (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        immVal = signExtend imm
    in WriteBackGPR rd (val1 .&. immVal)
  SLLI  (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        shamt = imm .&. 0x3F
    in WriteBackGPR rd (val1 `shiftL` fromIntegral shamt)
  SRLI  (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        shamt = imm .&. 0x3F
    in WriteBackGPR rd (val1 `shiftR` fromIntegral shamt)
  SRAI  (ITypeFields rd _ rs1 imm) ->
    let val1 = unpack (pack (extractRegVal rs1) :: BitVector 64) :: Signed 64
        shamt = imm .&. 0x3F
    in WriteBackGPR rd (bitCoerce (val1 `shiftR` fromIntegral shamt))
  SLTI  (ITypeFields rd _ rs1 imm) ->
    let val1 = unpack (pack (extractRegVal rs1) :: BitVector 64) :: Signed 64
        immVal = unpack (pack (signExtend imm) :: BitVector 64) :: Signed 64
    in WriteBackGPR rd (if val1 < immVal then 1 else 0)
  SLTIU (ITypeFields rd _ rs1 imm) ->
    let val1 = extractRegVal rs1
        immVal = signExtend imm
    in WriteBackGPR rd (if val1 < immVal then 1 else 0)
  LB    (ITypeFields rd _ rs1 imm) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm
    in Execute.ReadRequest (BusTypes.ReadRequest (baseAddr + offset) SizeByte) addr
  LH    (ITypeFields rd _ rs1 imm) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm
    in Execute.ReadRequest (BusTypes.ReadRequest (baseAddr + offset) SizeHalfWord) addr
  LW    (ITypeFields rd _ rs1 imm) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm
    in Execute.ReadRequest (BusTypes.ReadRequest (baseAddr + offset) SizeFullWord) addr
  LBU   (ITypeFields rd _ rs1 imm) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm
    in Execute.ReadRequest (BusTypes.ReadRequest (baseAddr + offset) SizeByte) addr
  LHU   (ITypeFields rd _ rs1 imm) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm
    in Execute.ReadRequest (BusTypes.ReadRequest (baseAddr + offset) SizeHalfWord) addr
  JALR  (ITypeFields rd _ rs1 imm) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm
        target = baseAddr + offset
    in if rd /= 0 then WriteBackGPR rd (addr + 4) else Jump target
  ECALL (ITypeFields _ _ _ _) ->
    Execute.DecodeException EnvironmentCallFromMMode addr  -- Assuming Machine mode for now
  EBREAK (ITypeFields _ _ _ _) ->
    Execute.DecodeException Breakpoint addr

  -- S-Type Instructions
  SB (STypeFields _ rs1 rs2 imm12) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm12
        val = extractRegVal rs2
    in Execute.WriteRequest (BusTypes.WriteRequest (baseAddr + offset) (BusByte (resize val))) addr
  SH (STypeFields _ rs1 rs2 imm12) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm12
        val = extractRegVal rs2
    in Execute.WriteRequest (BusTypes.WriteRequest (baseAddr + offset) (BusHalfWord (resize val))) addr
  SW (STypeFields _ rs1 rs2 imm12) ->
    let baseAddr = extractRegVal rs1
        offset = signExtend imm12
        val = extractRegVal rs2
    in Execute.WriteRequest (BusTypes.WriteRequest (baseAddr + offset) (BusFullWord (resize val))) addr

  -- B-Type Instructions
  BEQ (BTypeFields _ rs1 rs2 imm13) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
        offset = signExtend imm13
    in if val1 == val2 then Jump (addr + offset) else Jump (addr + 4)
  BNE (BTypeFields _ rs1 rs2 imm13) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
        offset = signExtend imm13
    in if val1 /= val2 then Jump (addr + offset) else Jump (addr + 4)
  BLT (BTypeFields _ rs1 rs2 imm13) ->
    let val1 = unpack (pack (extractRegVal rs1) :: BitVector 64) :: Signed 64
        val2 = unpack (pack (extractRegVal rs2) :: BitVector 64) :: Signed 64
        offset = signExtend imm13
    in if val1 < val2 then Jump (addr + offset) else Jump (addr + 4)
  BGE (BTypeFields _ rs1 rs2 imm13) ->
    let val1 = unpack (pack (extractRegVal rs1) :: BitVector 64) :: Signed 64
        val2 = unpack (pack (extractRegVal rs2) :: BitVector 64) :: Signed 64
        offset = signExtend imm13
    in if val1 >= val2 then Jump (addr + offset) else Jump (addr + 4)
  BLTU (BTypeFields _ rs1 rs2 imm13) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
        offset = signExtend imm13
    in if val1 < val2 then Jump (addr + offset) else Jump (addr + 4)
  BGEU (BTypeFields _ rs1 rs2 imm13) ->
    let val1 = extractRegVal rs1
        val2 = extractRegVal rs2
        offset = signExtend imm13
    in if val1 >= val2 then Jump (addr + offset) else Jump (addr + 4)

  -- U-Type Instructions
  LUI (UTypeFields rd imm20) ->
    let val = shiftL (resize imm20) 12
    in WriteBackGPR rd val
  AUIPC (UTypeFields rd imm20) ->
    let val = addr + shiftL (resize imm20) 12
    in WriteBackGPR rd val

  -- J-Type Instructions
  JAL (JTypeFields rd imm21) ->
    let offset = signExtend imm21
    in if rd /= 0 then WriteBackGPR rd (addr + 4) else Jump (addr + offset)

execute (Decode.DecodeException e addr) = Execute.DecodeException e addr
execute (Decode.InstructionException e addr) = Execute.InstructionException e addr