module Render.Debug.Model
  ( Model
  , VertexAttrs
  , vkVertexAttrs

  , InstanceAttrs(..)
  , instanceAttrs

  , StorableAttrs
  , storableAttrs1

  , InstanceBuffers(..)

  , TextureParams(..)
  , vkInstanceTexture

    -- TODO: extract and merge with UnlitTextured
  , allocateInstancesWith
  , allocateInstancesCoherent
  , allocateInstancesCoherent_
  , updateCoherentResize_
  , Transform
  ) where

import RIO


import Foreign (Storable(..))
import Geomancy (Transform, Vec2, Vec4)
import Geomancy.Vec3 qualified as Vec3
import RIO.Vector.Storable qualified as Storable
import UnliftIO.Resource (MonadResource, ReleaseKey, ResourceT, allocate)
import Vulkan.Core10 qualified as Vk
import Vulkan.NamedType ((:::))
import Vulkan.Zero (Zero(..))

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

type Model buf = Model.Indexed buf Vec3.Packed VertexAttrs

type VertexAttrs = "uv" ::: Vec2

vkVertexAttrs :: [Vk.Format]
vkVertexAttrs =
  [ Vk.FORMAT_R32G32_SFLOAT -- vTexCoord :: vec2
  ]

-- | Data for a single element.
data InstanceAttrs = InstanceAttrs
  { textureParams :: TextureParams
  , transformMat4 :: Transform
  }

instance Zero InstanceAttrs where
  zero = InstanceAttrs
    { textureParams = zero
    , transformMat4 = mempty
    }

instanceAttrs :: Int32 -> Int32 -> [Transform] -> InstanceAttrs
instanceAttrs samplerId textureId transforms = InstanceAttrs
  { textureParams = zero
      { tpSamplerId = samplerId
      , tpTextureId = textureId
      }
  , transformMat4 = mconcat transforms
  }

-- | Intermediate data to be shipped.
type StorableAttrs =
  ( Storable.Vector TextureParams
  , Storable.Vector Transform
  )

storableAttrs1 :: Int32 -> Int32 -> [Transform] -> StorableAttrs
storableAttrs1 samplerId textureId transforms =
  ( Storable.singleton textureParams
  , Storable.singleton transformMat4
  )
  where
    InstanceAttrs{..} = instanceAttrs
      samplerId
      textureId
      transforms

-- | GPU-bound data.
data InstanceBuffers textureStage transformStage = InstanceBuffers
  { ibTexture   :: InstanceTexture textureStage
  , ibTransform :: InstanceTransform transformStage
  }

type InstanceTexture stage = Buffer.Allocated stage TextureParams

type InstanceTransform stage = Buffer.Allocated stage Transform

instance Model.HasVertexBuffers (InstanceBuffers textureStage transformStage) where
  type VertexBuffersOf (InstanceBuffers textureStage transformStage) = InstanceAttrs

  {-# INLINE getVertexBuffers #-}
  getVertexBuffers InstanceBuffers{..} =
    [ Buffer.aBuffer ibTexture
    , Buffer.aBuffer ibTransform
    ]

  {-# INLINE getInstanceCount #-}
  getInstanceCount InstanceBuffers{..} =
    min
      (Buffer.aUsed ibTexture)
      (Buffer.aUsed ibTransform)

data TextureParams = TextureParams
  { tpScale     :: Vec2
  , tpOffset    :: Vec2
  , tpGamma     :: Vec4
  , tpSamplerId :: Int32
  , tpTextureId :: Int32
  }
  deriving (Show)

instance Zero TextureParams where
  zero = TextureParams
    { tpScale     = 1
    , tpOffset    = 0
    , tpGamma     = 1.0
    , tpSamplerId = minBound
    , tpTextureId = minBound
    }

instance Storable TextureParams where
  alignment ~_ = 4

  sizeOf ~_ = 8 + 8 + 16 + 4 + 4

  poke ptr TextureParams{..} = do
    pokeByteOff ptr  0 tpScale
    pokeByteOff ptr  8 tpOffset
    pokeByteOff ptr 16 tpGamma
    pokeByteOff ptr 32 tpSamplerId
    pokeByteOff ptr 36 tpTextureId

  peek ptr = do
    tpScale     <- peekByteOff ptr  0
    tpOffset    <- peekByteOff ptr  8
    tpGamma     <- peekByteOff ptr 16
    tpSamplerId <- peekByteOff ptr 32
    tpTextureId <- peekByteOff ptr 36
    pure TextureParams{..}

vkInstanceTexture :: [Vk.Format]
vkInstanceTexture =
  [ Vk.FORMAT_R32G32B32A32_SFLOAT -- iTextureScaleOffset :: vec4
  , Vk.FORMAT_R32G32B32A32_SFLOAT -- iTextureGamma       :: vec4
  , Vk.FORMAT_R32G32_SINT         -- iTextureIds         :: ivec2
  ]

allocateInstancesWith
  :: ( MonadResource m
     , MonadUnliftIO m
     )
  => (Vk.BufferUsageFlagBits -> Int -> Storable.Vector TextureParams -> m (InstanceTexture texture))
  -> (Vk.BufferUsageFlagBits -> Int -> Storable.Vector Transform -> m (InstanceTransform transform))
  -> (forall stage a . Buffer.Allocated stage a -> m ())
  -> [InstanceAttrs]
  -> m (ReleaseKey, InstanceBuffers texture transform)
allocateInstancesWith createTextures createTransforms bufferDestroy instances = do
  ul <- askUnliftIO
  allocate (create ul) (destroy ul)
  where
    textures   = Storable.fromList $ map textureParams instances
    transforms = Storable.fromList $ map transformMat4 instances
    numInstances = Storable.length textures

    create (UnliftIO ul) = ul do
      ibTexture   <- createTextures   Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT numInstances textures
      ibTransform <- createTransforms Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT numInstances transforms
      pure InstanceBuffers{..}

    destroy (UnliftIO ul) InstanceBuffers{..} = ul do
      bufferDestroy ibTexture
      bufferDestroy ibTransform

allocateInstancesCoherent
  :: ( MonadReader env m
     , HasVulkan env
     , MonadResource m
     , MonadUnliftIO m
     )
  => [InstanceAttrs]
  -> m (ReleaseKey, InstanceBuffers 'Buffer.Coherent 'Buffer.Coherent)
allocateInstancesCoherent instances = do
  context <- ask
  allocateInstancesWith
    (Buffer.createCoherent context)
    (Buffer.createCoherent context)
    (Buffer.destroy context)
    instances

allocateInstancesCoherent_
  :: (HasVulkan env)
  => Int
  -> ResourceT (RIO env) (InstanceBuffers 'Buffer.Coherent 'Buffer.Coherent)
allocateInstancesCoherent_ n =
  fmap snd $ allocateInstancesCoherent (replicate n zero)

updateCoherentResize_
  :: ( HasVulkan context
     , MonadUnliftIO m
     )
  => context
  -> InstanceBuffers 'Buffer.Coherent 'Buffer.Coherent
  -> (Storable.Vector TextureParams, Storable.Vector Transform)
  -> m (InstanceBuffers 'Buffer.Coherent 'Buffer.Coherent)
updateCoherentResize_ context InstanceBuffers{..} (textures, transforms) =
  InstanceBuffers
    <$> Buffer.updateCoherentResize_ context ibTexture textures
    <*> Buffer.updateCoherentResize_ context ibTransform transforms