module Render.Lit.Textured.Model
  ( Model
  , VertexAttrs(..)
  , vkVertexAttrs

  , InstanceAttrs(..)
  , InstanceBuffers(..)

  , TextureParams(..)
  , vkInstanceTexture

  , allocateInstancesWith
  , Transform
  ) where

import RIO


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

import Resource.Buffer qualified as Buffer
import Resource.Model qualified as Model

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

data VertexAttrs = VertexAttrs
  { vaTexCoord :: Vec2
  , vaNormal   :: Vec3.Packed
  }
  deriving (Eq, Ord, Show, Generic)

instance Storable VertexAttrs where
  alignment ~_ = 4

  sizeOf ~_ = 8 + 12

  peek ptr = do
    vaTexCoord <- peekByteOff ptr 0
    vaNormal   <- peekByteOff ptr 8
    pure VertexAttrs{..}

  poke ptr VertexAttrs{..} = do
    pokeByteOff ptr 0 vaTexCoord
    pokeByteOff ptr 8 vaNormal

vkVertexAttrs :: [Vk.Format]
vkVertexAttrs =
  [ Vk.FORMAT_R32G32_SFLOAT    -- vTexCoord :: vec2
  , Vk.FORMAT_R32G32B32_SFLOAT -- vNormal   :: vec3
  ]

data InstanceAttrs = InstanceAttrs
  { textureParams :: TextureParams
  , transformMat4 :: Transform
  }

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

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 ~_ = 8

  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 -> VectorS.Vector TextureParams -> m (InstanceTexture texture))
  -> (Vk.BufferUsageFlagBits -> Int -> VectorS.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   = VectorS.fromList $ map textureParams instances
    transforms = VectorS.fromList $ map transformMat4 instances
    numInstances = VectorS.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