module Resource.Mesh.Codec where

import RIO

import Codec.Compression.Zstd qualified as Zstd
import Codec.Serialise qualified as CBOR
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.Typeable (typeRep, typeRepTyCon)
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 Engine.Vulkan.Types (MonadVulkan, Queues)
import Resource.Buffer qualified as Buffer
import Resource.Model qualified as Model
import Resource.Region qualified as Region

-- * Format meta

pattern VER_BREAKS :: Word8
pattern VER_BREAKS = 3

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
     , CBOR.Serialise 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

  let (metSize, metDigest, metCompressed) = encodeCBOR 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 metSize
      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)

encodeCBOR :: CBOR.Serialise a => a -> (Int, ByteString, ByteString)
encodeCBOR stuff =
  ( ByteString.length buf
  , MD5.hash buf
  , Zstd.compress Zstd.maxCLevel buf
  )
  where
    buf = BSL.toStrict (CBOR.serialise stuff)

-- * Decoding

loadIndexed
  :: ( Storable attrs
     , Storable nodes
     , CBOR.Serialise meta
     , Show meta
     , Typeable nodes
     , HasLogFunc env
     , MonadResource m
     , MonadVulkan env m
     )
  => Queues Vk.CommandPool
  -> FilePath
  -> m
    ( 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

  Region.run do
    logDebug $ "Staging " <> fromString fp
    indexed <- Model.createStaged (Just $ fromString fp) pools positions attrs indices
    Model.registerIndexed_ indexed
    pure (meta, nodes, indexed)

loadBlobs
  :: forall attrs env nodes meta m
  .  ( Storable attrs
     , CBOR.Serialise meta
     , Storable nodes
     , Typeable nodes
     , HasLogFunc env
     , MonadReader env m
     , MonadIO m
     )
  => FilePath
  -> m
    ( 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      <- decodeCBOR  "Metadata"   metDigest sizeOfMeta          zMetas

      pure (verTweaks, 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

decodeCBOR
  :: ( CBOR.Serialise a
     , MonadFail m
     )
  => String
  -> ByteString
  -> Int
  -> ByteString
  -> m a
decodeCBOR label digest expectedSize zBytes =
  case Zstd.decompress zBytes of
    Zstd.Error err ->
      fail $ label <> ": zstd error (" <> err <> ")"
    Zstd.Skip ->
      fail $ label <> ": empty zstd"
    Zstd.Decompress bytes -> do
      guardEq (label <> " size") expectedSize (ByteString.length bytes)
      guardEq (label <> " hash") digest (MD5.hash bytes)
      case CBOR.deserialiseOrFail (BSL.fromStrict bytes) of
        Right value ->
          pure value
        Left (CBOR.DeserialiseFailure _off err) ->
          fail $ unlines
            [ label <> " deserialise failure:"
            , err
            ]

-- * 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
      ]