KeyValue.hs

module Codec.Ktx.KeyValue
  ( KeyValueData

    -- * Writing
  , Value(..)
  , text
  , bytes
  , number

    -- * Reading
  , FromValue(..)
  , textual

    -- * Binary operations
  , getDataLe
  , getData

  , putDataLe
  , putData
  ) where

import Data.Binary.Get (Get, getWord32le, getByteString, isolate, skip)
import Data.Binary.Put (Put, putByteString, putWord32le)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Foldable (for_)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Text qualified as Text (pack)
import Data.Text.Encoding qualified as Text
import Data.Word (Word32)
import GHC.Generics (Generic)

type KeyValueData = Map Text Value

{- | A wrapper for raw data.

Use "FromValue"/"ToValue" to process.
-}
newtype Value = Value ByteString
  deriving (Eq, Show, Generic)

class FromValue a where
  fromValue :: Value -> Maybe a

instance FromValue Text where
  fromValue (Value bs)
    | BS.null bs =
        Nothing
    | BS.last bs == 0x00 =
        either (const Nothing) Just $
          Text.decodeUtf8' $ BS.init bs
    | otherwise =
        Nothing

instance FromValue ByteString where
  fromValue (Value bs) = Just bs

text :: Text -> Value
text t = Value $ BS.snoc (Text.encodeUtf8 t) 0x00

bytes :: ByteString -> Value
bytes = Value

number :: (Num a, Show a) => a -> Value
number = text . Text.pack . show

-- | Extract all valid (null-terminated utf8) values.
textual :: KeyValueData -> Map Text Text
textual = Map.mapMaybe fromValue

{-# INLINE getDataLe #-}
getDataLe :: Int -> Get KeyValueData
getDataLe = getData getWord32le

getData :: Get Word32 -> Int -> Get KeyValueData
getData getSize bytesOfKeyValueData =
  isolate bytesOfKeyValueData $
    go bytesOfKeyValueData []
  where
    go remains acc
      | remains == 0 =
          pure $ Map.fromList acc

      | remains < 0 =
          fail "Attempted to read beyond bytesOfKeyValueData"

      | otherwise = do
          keyAndValueByteSize <- fmap fromIntegral getSize
          let paddingSize = 3 - ((keyAndValueByteSize + 3) `rem` 4)

          keyAndValue <- getByteString keyAndValueByteSize
          skip paddingSize

          {- XXX: Spec says:
              Any byte value is allowed.
              It is encouraged that the value be a NUL terminated UTF-8 string but this is not required.
              If the Value data is a string of bytes then the NUL termination
              should be included in the keyAndValueByteSize byte count
              (but programs that read KTX files must not rely on this).
          -}
          let
            (keyBS, valueBS) = BS.span (/= 0x00) keyAndValue
            key = Text.decodeUtf8 keyBS
            value = Value $ BS.drop 1 valueBS

          go
            (remains - keyAndValueByteSize - 4 - paddingSize)
            ((key, value) : acc)

{-# INLINE putDataLe #-}
putDataLe :: KeyValueData -> Put
putDataLe = putData putWord32le

putData :: (Word32 -> Put) -> KeyValueData -> Put
putData putSize kvs =
  for_ (Map.toList kvs) \(key, Value value) -> do
    let
      keyAndValue = mconcat [Text.encodeUtf8 key, BS.singleton 0x00, value]
      keyAndValueByteSize = BS.length keyAndValue
      paddingSize = 3 - ((keyAndValueByteSize + 3) `rem` 4)

    putSize (fromIntegral keyAndValueByteSize)
    putByteString keyAndValue
    putByteString $ BS.replicate paddingSize 0