Codec.hs

module Resource.Mesh.Codec where

import RIO

import Codec.Compression.Zstd qualified as Zstd
import Crypto.Hash.MD5 qualified as MD5
import Data.Binary.Get (runGet)
import Data.Binary.Get qualified as Get
import Data.Binary.Put (runPut)
import Data.Binary.Put qualified as Put
import Data.ByteString.Internal qualified as BSI
import Data.ByteString.Unsafe (unsafePackCStringLen)
import Data.Vector qualified as Vector
import Data.Vector.Generic qualified as Generic
import Data.Vector.Storable qualified as Storable
import Foreign qualified
import Geomancy.Vec3 qualified as Vec3
import RIO.ByteString qualified as ByteString
import RIO.ByteString.Lazy qualified as BSL
import System.IO.Unsafe (unsafePerformIO)
import UnliftIO.Resource (MonadResource)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Data.Typeable

import Engine.Vulkan.Types (HasVulkan, Queues)
import Resource.Buffer qualified as Buffer
import Resource.Model qualified as Model

-- * Format meta

pattern VER_BREAKS :: Word8
pattern VER_BREAKS = 2

pattern VER_TWEAKS :: Word8
pattern VER_TWEAKS = 0

-- * Encoding

encodeFile
  :: forall vp vi va vn attrs nodes meta env
  .  ( Generic.Vector vp Vec3.Packed
     , Generic.Vector vi Word32
     , Generic.Vector va attrs
     , Generic.Vector vn nodes
     , Storable attrs
     , Storable nodes
     , Storable meta
     , HasLogFunc env
     )
  => FilePath
  -> vp Vec3.Packed
  -> vi Word32
  -> va attrs
  -> vn nodes
  -> meta
  -> RIO env ()
encodeFile fp positions indices attrs nodes meta = do
  (posDigest, posCompressed) <- encodeItems positions
  logDebug $ "Position digest: " <> displayShow posDigest

  (indDigest, indCompressed) <- encodeItems indices
  logDebug $ "Index digest: " <> displayShow indDigest

  (attDigest, attCompressed) <- encodeItems attrs
  logDebug $ "Attribute digest: " <> displayShow attDigest

  (nodDigest, nodCompressed) <- encodeItems nodes
  logDebug $ "Node digest: " <> displayShow nodDigest

  (metDigest, metCompressed) <- encodeItems $ Storable.singleton meta
  logDebug $ "Meta digest: " <> displayShow metDigest

  withFile fp WriteMode \out -> do
    BSL.hPut out $ runPut do
      -- 0x00 + 4 + 4
      Put.putStringUtf8 "🌋📦"

      -- 0x08 + 1 + 1
      Put.putWord8 VER_BREAKS
      Put.putWord8 VER_TWEAKS

      -- 0x0A + 2 + 4
      -- XXX: reserved
      Put.putWord16le maxBound
      Put.putWord32le maxBound

      -- 0x10 + 4 + 4
      Put.putWord32le . fromIntegral $ Generic.length positions
      Put.putWord32le . fromIntegral $ ByteString.length posCompressed

      -- 0x18 + 4 + 4
      Put.putWord32le . fromIntegral $ Generic.length indices
      Put.putWord32le . fromIntegral $ ByteString.length indCompressed

      -- 0x20 + 4 + 4
      Put.putWord32le . fromIntegral $ Foreign.sizeOf (error "sizeOf" :: attrs)
      Put.putWord32le . fromIntegral $ ByteString.length attCompressed

      -- 0x28 + 4 + 4
      Put.putWord32le . fromIntegral $ Foreign.sizeOf (error "sizeOf" :: nodes)
      Put.putWord32le . fromIntegral $ ByteString.length nodCompressed

      -- 0x30 + 4 + 4
      Put.putWord32le . fromIntegral $ Foreign.sizeOf (error "sizeOf" :: meta)
      Put.putWord32le . fromIntegral $ ByteString.length metCompressed

      -- 0x38 + 4 + 4
      -- XXX: reserved
      Put.putWord32le maxBound
      Put.putWord32le maxBound

      -- 0x40 + 16 + 16 + 16 + 16 + 16
      Put.putByteString posDigest
      Put.putByteString indDigest
      Put.putByteString attDigest
      Put.putByteString nodDigest
      Put.putByteString metDigest

      -- 0x90 + posCompressed
      Put.putByteString posCompressed

      -- 0x90 + posCompressed + indCompressed
      Put.putByteString indCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed
      Put.putByteString attCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed + nodCompressed
      Put.putByteString nodCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed + nodCompressed + metCompressed
      Put.putByteString metCompressed

      -- 0x90 + posCompressed + indCompressed + attCompressed + nodCompressed + metCompressed + 4 + 4
      Put.putStringUtf8 "📦🌋"

encodeItems
  :: ( Storable a
     , Generic.Vector v a
     , MonadIO m
     )
  => v a
  -> m (ByteString, ByteString)
encodeItems items = do
  let bytes = Storable.unsafeCast @_ @Word8 $ Vector.convert items
  liftIO $ Storable.unsafeWith bytes \ptr -> do
    buf <- unsafePackCStringLen
      ( Foreign.castPtr ptr
      , Storable.length bytes
      )
    let
      -- XXX: Process buffer before bytes/ptr go out of scope!
      !bufHash = MD5.hash buf
      !zBuf = Zstd.compress Zstd.maxCLevel buf
    pure (bufHash, zBuf)

-- * Decoding

loadIndexed
  :: ( Storable attrs
     , Storable nodes
     , Storable meta
     , Show meta
     , Typeable nodes
     , HasVulkan env
     , HasLogFunc env
     , MonadResource (RIO env)
     )
  => Queues Vk.CommandPool
  -> FilePath
  -> RIO env
    ( Resource.ReleaseKey
    , (meta, Storable.Vector nodes, Model.Indexed 'Buffer.Staged Vec3.Packed attrs)
    )
loadIndexed pools fp = do
  logInfo $ "Loading " <> fromString fp
  (meta, nodes, (positions, indices, attrs)) <- loadBlobs fp
  logDebug $ displayShow meta

  logDebug $ "Staging " <> fromString fp
  context <- ask
  indexed <- Model.createStaged context pools positions attrs indices
  key <- Resource.register $ Model.destroyIndexed context indexed
  pure (key, (meta, nodes, indexed))

loadBlobs
  :: forall attrs env nodes meta
  .  ( Storable attrs
     , Storable meta
     , Storable nodes
     , Typeable nodes
     , HasLogFunc env
     )
  => FilePath
  -> RIO env
    ( meta
    , Storable.Vector nodes
    , ( Storable.Vector Vec3.Packed
      , Storable.Vector Word32
      , Storable.Vector attrs
      )
    )
loadBlobs fp = do
  blob <- BSL.readFile fp
  let
    getter = do
      magicStart <- Get.getByteString (4+4)

      verBreaks <- Get.getWord8
      verTweaks <- Get.getWord8
      guardEq "Codec version" VER_BREAKS verBreaks

      _reserved16 <- Get.getWord16le
      _reserved32 <- Get.getWord32le

      numPositions <- fmap fromIntegral Get.getWord32le
      lenPositions <- fmap fromIntegral Get.getWord32le

      numIndices <- fmap fromIntegral Get.getWord32le
      lenIndices <- fmap fromIntegral Get.getWord32le

      sizeOfAttr <- fmap fromIntegral Get.getWord32le
      lenAttrs   <- fmap fromIntegral Get.getWord32le
      guardEq "Attribute size" (Foreign.sizeOf (error "sizeOfAttr" :: attrs)) sizeOfAttr

      sizeOfNode <- fmap fromIntegral Get.getWord32le
      lenNodes   <- fmap fromIntegral Get.getWord32le
      guardEq
        ("Node size for " <> show (typeRepTyCon . typeRep $ Proxy @nodes))
        (Foreign.sizeOf (error "sizeOfNode" :: nodes))
        sizeOfNode

      sizeOfMeta <- fmap fromIntegral Get.getWord32le
      lenMeta    <- fmap fromIntegral Get.getWord32le
      guardEq "Meta size" (Foreign.sizeOf (error "sizeOf" :: meta)) sizeOfMeta

      _reserved32 <- Get.getWord32le
      _reserved32 <- Get.getWord32le

      posDigest <- Get.getByteString 16
      indDigest <- Get.getByteString 16
      attDigest <- Get.getByteString 16
      nodDigest <- Get.getByteString 16
      metDigest <- Get.getByteString 16

      -- XXX: end for static part for VER_BREAKS
      staticDone <- Get.bytesRead
      guardEq ("End of static part for v" <> show VER_BREAKS) 0x90 staticDone

      let payloadSize = fromIntegral (lenPositions + lenIndices + lenAttrs + lenNodes + lenMeta)
      guardEq "Blob size" (BSL.length blob) $
        staticDone + payloadSize + 4 + 4

      zPositions <- Get.getByteString lenPositions
      zIndices   <- Get.getByteString lenIndices
      zAttrs     <- Get.getByteString lenAttrs
      zNodes     <- Get.getByteString lenNodes
      zMetas     <- Get.getByteString lenMeta

      magicFinish <- Get.getByteString (4+4)

      let magicReverse = ByteString.drop 4 magicFinish <> ByteString.take 4 magicFinish
      guardEq "Magic final" magicStart magicReverse

      positions <- decodeItems "Positions"  posDigest (Just numPositions) zPositions
      indices   <- decodeItems "Indices"    indDigest (Just numIndices)   zIndices
      attrs     <- decodeItems "Attributes" attDigest (Just numPositions) zAttrs
      nodes     <- decodeItems "Nodes"      nodDigest Nothing             zNodes
      meta      <- decodeItems "Metadata"   metDigest (Just 1)            zMetas

      pure (verTweaks, Storable.head meta, nodes, (positions, indices, attrs))

  let (verTweaks, meta, nodes, blobs) = runGet getter blob

  when (verTweaks /= VER_TWEAKS) $
    logWarn $ mconcat
      [ "Format tweak version mismatch: "
      , display verTweaks
      , " /= "
      , display VER_TWEAKS
      ]

  pure (meta, nodes, blobs)

decodeItems
  :: forall item m
  . (Storable item, MonadFail m)
  => String
  -> ByteString
  -> Maybe Int
  -> ByteString
  -> m (Storable.Vector item)
decodeItems label digest expectedSize zBytes =
  case Zstd.decompress zBytes of
    Zstd.Error err ->
      fail err
    Zstd.Skip -> do
      guardEq (label <> " size") 0 itemSize
      case expectedSize of
        Just size ->
          pure . Storable.replicate size $
            unsafePerformIO (Foreign.peek Foreign.nullPtr)
        Nothing ->
          pure Storable.empty
    Zstd.Decompress bytes -> do
      let (buf, bufOff, bufLen) = BSI.toForeignPtr bytes
      guardEq (label <> " buffer offset") 0 bufOff
      case expectedSize of
        Nothing ->
          pure ()
        Just size ->
          guardEq (label <> " buffer size") (itemSize * size) bufLen
      guardEq (label <> " hash") digest (MD5.hash bytes)
      let
        !items =
          Storable.unsafeCast @Word8 @item $
            Storable.unsafeFromForeignPtr0 buf bufLen
      case expectedSize of
        Nothing ->
          pure ()
        Just size ->
          guardEq (label <> " size") size (Storable.length items)
      pure items
  where
    itemSize = Foreign.sizeOf @item undefined

-- * Utils

guardEq :: (MonadFail m, Show a, Eq a) => String -> a -> a -> m ()
guardEq label expected got = do
  if expected == got then
    -- traceM . fromString $ label <> " match: " <> show got
    pure ()
  else
    fail $ unlines
      [ label <> " mismatch"
      , "\tExpected: " <> show expected
      , "\tGot     : " <> show got
      ]