{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

module Peripherals.Ram(
  initRamFromFile,
  RamAddr,
  Ram,
  RamLine,
  bytesInRam,
  read,
  write,
  ) where

import Clash.Prelude hiding (empty, read)
import qualified Prelude as P
import qualified Data.ByteString.Lazy as BL
import Data.Binary.Get
import Data.Int (Int32)
import qualified Clash.Sized.Vector as Vec
import Types(Addr, FullWord, DoubleWord)
import BusTypes(
    TransactionSize(..),
    BusVal(..),
  )

-- vector depth has to be known statically at compile time
#ifndef _RAM_DEPTH
#define _RAM_DEPTH 1024
#endif

-- TODO : replace Unsigned 32 with BusVal types later...
type Ram = Vec _RAM_DEPTH (Unsigned 32)
type RamAddr = Unsigned (CLog 2 _RAM_DEPTH)
type RamLine = Unsigned 32
bytesInRam :: Addr
bytesInRam = _RAM_DEPTH * 4

read :: TransactionSize -> RamAddr -> Ram -> BusVal
read SizeByte addr ram = BusByte $ unpack byte
  where
    word = ram !! addr
    byteOffset :: BitVector 2
    byteOffset = slice d1 d0 addr
    byte = case byteOffset of
      0b00 -> slice d31 d24 word
      0b01 -> slice d23 d16 word
      0b10 -> slice d15 d8 word
      0b11 -> slice d7 d0 word

read SizeHalfWord addr ram = BusHalfWord $ unpack halfWord
  where
    word = ram !! addr
    halfWordOffset :: Unsigned 1
    halfWordOffset = unpack $ slice d0 d0 addr
    halfWord = case halfWordOffset of
      0b0 -> slice d31 d16 word
      0b1 -> slice d15 d0 word

read SizeFullWord addr ram = BusFullWord fullWord
  where
    fullWord = ram !! addr

read SizeDoubleWord addr ram = BusDoubleWord doubleWord
  where
    doubleWord = bitCoerce $ bitCoerce word0 ++# bitCoerce word1
    word0 = readFullWordHelper ram addr
    word1 = readFullWordHelper ram (addr + 1)

read SizeQuadWord addr ram = BusQuadWord quadWord
  where
    quadWord = bitCoerce $ bitCoerce dword0 ++# bitCoerce dword1
    dword0 = readDoubleWordHelper ram addr
    dword1 = readDoubleWordHelper ram (addr + 2)

readFullWordHelper :: Ram -> RamAddr -> FullWord
readFullWordHelper ram addr = ram !! addr

readDoubleWordHelper :: Ram -> RamAddr -> DoubleWord
readDoubleWordHelper ram addr = bitCoerce $ bitCoerce word0 ++# bitCoerce word1
  where
    word0 = readFullWordHelper ram addr
    word1 = readFullWordHelper ram (addr + 1)

write :: BusVal -> RamAddr -> Ram -> Ram
write (BusByte byte) addr ram = replace addr updatedWord ram
  where
    word = ram !! addr
    byteOffset :: BitVector 2
    byteOffset = slice d1 d0 addr
    updatedWord = case byteOffset of
      0b00 -> setSlice d31 d24 (pack byte) word
      0b01 -> setSlice d23 d16 (pack byte) word
      0b10 -> setSlice d15 d8  (pack byte) word
      0b11 -> setSlice d7 d0   (pack byte) word

write (BusHalfWord halfWord) addr ram = replace addr updatedWord ram
  where
    word = ram !! addr
    halfWordOffset :: Unsigned 1
    halfWordOffset = unpack $ slice d0 d0 addr
    updatedWord = case halfWordOffset of
      0b0 -> setSlice d31 d16 (pack halfWord) word
      0b1 -> setSlice d15 d0  (pack halfWord) word

write (BusFullWord fullWord) addr ram = replace addr fullWord ram

write (BusDoubleWord doubleWord) addr ram = ram''
  where
    (word0, word1) = bitCoerce doubleWord
    ram' = replace addr word0 ram
    ram'' = replace (addr + 1) word1 ram'

write (BusQuadWord quadWord) addr ram = ram''''
  where
    (dword0 :: DoubleWord, dword1 :: DoubleWord) =
      bitCoerce quadWord

    (word0 :: FullWord, word1 :: FullWord) =
      bitCoerce dword0

    (word2 :: FullWord, word3 :: FullWord) =
      bitCoerce dword1

    ram' = replace addr word0 ram
    ram'' = replace (addr + 1) word1 ram'
    ram''' = replace (addr + 2) word2 ram''
    ram'''' = replace (addr + 3) word3 ram'''

initRamFromFile :: FilePath -> IO (Maybe Ram)
initRamFromFile filePath =
    let
      initRam = Vec.replicate (SNat :: SNat _RAM_DEPTH) 0
    in
      do
      bs <- readFileIntoByteString filePath
      let ints = getInts bs
      pure $ populateVectorFromInt32 ints initRam

readFileIntoByteString :: FilePath -> IO BL.ByteString
readFileIntoByteString filePath = BL.readFile filePath

-- Define a function to read a ByteString and convert to [Int32]
getInts :: BL.ByteString -> [Int32]
getInts bs = runGet listOfInts bs
  where
    listOfInts = do
      empty <- isEmpty
      if empty
        then pure []
      else do
        i <- getInt32le    -- Parse a single Int32 from the stream
        rest <- listOfInts -- Recursively parse the rest
        pure (i : rest)

-- Adjusts the length of a list of integers by either truncating or padding with zeros
populateVectorFromInt32 ::
    KnownNat n =>
    [Int32] ->
    Vec n (Unsigned 32) ->
    Maybe (Vec n (Unsigned 32))
populateVectorFromInt32 ls v = Vec.fromList adjustedLs
  where
    vecLen = length v
    adjustedLs = fromIntegral <$> adjustLength vecLen ls
    adjustLength :: Int -> [Int32] -> [Int32]
    adjustLength n xs = P.take n (xs P.++ P.repeat 0)