module Rattletrap.Primitive.CompressedWord where

import qualified Data.Binary.Bits.Get as BinaryBit
import qualified Data.Binary.Bits.Put as BinaryBit
import qualified Data.Bits as Bits

data CompressedWord = CompressedWord
  { compressedWordLimit :: Word
  , compressedWordValue :: Word
  } deriving (Eq, Ord, Show)

getCompressedWord :: Word -> BinaryBit.BitGet CompressedWord
getCompressedWord limit = do
  value <- getCompressedWordStep limit (getMaxBits limit) 0 0
  pure (CompressedWord limit value)

putCompressedWord :: CompressedWord -> BinaryBit.BitPut ()
putCompressedWord compressedWord = do
  let limit = compressedWordLimit compressedWord
  let value = compressedWordValue compressedWord
  let maxBits = getMaxBits limit
  let go position soFar =
        if position < maxBits
          then do
            let x = Bits.shiftL 1 position
            if maxBits > 1 && position == maxBits - 1 && soFar + x > limit
              then pure ()
              else do
                let bit = Bits.testBit value position
                BinaryBit.putBool bit
                let delta =
                      if bit
                        then x
                        else 0
                go (position + 1) (soFar + delta)
          else pure ()
  go 0 0

getMaxBits :: (Integral a, Integral b) => a -> b
getMaxBits x = do
  let n = max 1 (ceiling (logBase (2 :: Double) (fromIntegral (max 1 x))))
  if x < 1024 && x == 2 ^ n
    then n + 1
    else n

getCompressedWordStep :: Word -> Word -> Word -> Word -> BinaryBit.BitGet Word
getCompressedWordStep limit maxBits position value = do
  let x = Bits.shiftL 1 (fromIntegral position)
  if position < maxBits && value + x <= limit
    then do
      bit <- BinaryBit.getBool
      let newValue =
            if bit
              then value + x
              else value
      getCompressedWordStep limit maxBits (position + 1) newValue
    else pure value