module Rattletrap.Type.CompressedWord where

import qualified Data.Bits as Bits
import qualified Rattletrap.BitGet as BitGet
import qualified Rattletrap.BitPut as BitPut
import qualified Rattletrap.Schema as Schema
import qualified Rattletrap.Utility.Json as Json

-- | Although there's no guarantee that these values will not overflow, it's
-- exceptionally unlikely. Most 'CompressedWord's are very small.
data CompressedWord = CompressedWord
  { limit :: Word,
    value :: Word
  }
  deriving (Eq, Ord, Show)

instance Json.FromJSON CompressedWord where
  parseJSON = Json.withObject "CompressedWord" $ \object -> do
    limit <- Json.required object "limit"
    value <- Json.required object "value"
    pure CompressedWord {limit, value}

instance Json.ToJSON CompressedWord where
  toJSON x =
    Json.object [Json.pair "limit" $ limit x, Json.pair "value" $ value x]

schema :: Schema.Schema
schema =
  Schema.named "compressedWord" $
    Schema.object
      [ (Json.pair "limit" $ Json.object [Json.pair "type" "integer"], True),
        (Json.pair "value" $ Json.object [Json.pair "type" "integer"], True)
      ]

bitPut :: CompressedWord -> BitPut.BitPut
bitPut compressedWord =
  let limit_ = limit compressedWord
      value_ = value compressedWord
      maxBits = getMaxBits limit_
   in putCompressedWordStep limit_ value_ maxBits 0 0

putCompressedWordStep :: Word -> Word -> Int -> Int -> Word -> BitPut.BitPut
putCompressedWordStep limit_ value_ maxBits position soFar =
  if position < maxBits
    then do
      let x = Bits.shiftL 1 position :: Word
      if maxBits > 1 && position == maxBits - 1 && soFar + x > limit_
        then mempty
        else
          let bit = Bits.testBit value_ position
              delta = if bit then x else 0
           in BitPut.bool bit
                <> putCompressedWordStep
                  limit_
                  value_
                  maxBits
                  (position + 1)
                  (soFar + delta)
    else mempty

getMaxBits :: Word -> Int
getMaxBits x =
  let n :: Int
      n = max 1 (ceiling (logBase (2 :: Double) (fromIntegral (max 1 x))))
   in if x < 1024 && x == 2 ^ n then n + 1 else n

bitGet :: Word -> BitGet.BitGet CompressedWord
bitGet = bitGetNew

bitGetNew :: Word -> BitGet.BitGet CompressedWord
bitGetNew limit = do
  value <-
    if limit < 1
      then pure 0
      else do
        let numBits =
              max (0 :: Int)
                . subtract 1
                . ceiling
                . logBase (2 :: Double)
                $ fromIntegral limit
        partial <- BitGet.bits numBits
        let next = partial + Bits.shiftL 1 numBits
        if next > limit
          then pure partial
          else do
            x <- BitGet.bool
            pure $ if x then next else partial
  pure CompressedWord {limit, value}