module Render.Unlit.TileMap.Model
  ( Model
  , VertexAttrs
  , vkVertexAttrs

  , InstanceAttrs(..)
  -- , instanceAttrs

  , StorableAttrs
  -- , storableAttrs1

  , InstanceBuffers(..)

  , TileMapParams(..)
  , vkInstanceTileMap

  , allocateInstancesWith
  , allocateInstancesCoherent
  , allocateInstancesCoherent_
  , updateCoherentResize_
  , Transform
  ) where

import RIO

import Foreign (Storable(..))
import Geomancy (IVec4, Transform, Vec2)
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
  { tilemapParams :: TileMapParams
  , transformMat4 :: Transform
  }

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

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

-- | GPU-bound data.
data InstanceBuffers tilemapStage transformStage = InstanceBuffers
  { ibTileMap   :: InstanceTileMap tilemapStage
  , ibTransform :: InstanceTransform transformStage
  }

type InstanceTileMap stage = Buffer.Allocated stage TileMapParams

type InstanceTransform stage = Buffer.Allocated stage Transform

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

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

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

data TileMapParams = TileMapParams
  { tmpTextureIds         :: IVec4
  , tmpViewOffset         :: Vec2
  , tmpViewportSize       :: Vec2
  , tmpMapTextureSize     :: Vec2
  , tmpTilesetTextureSize :: Vec2
  , tmpTileSize           :: Vec2
  }
  deriving (Show)

instance Zero TileMapParams where
  zero = TileMapParams
    { tmpTextureIds         = 0
    , tmpViewOffset         = 0
    , tmpViewportSize       = 1
    , tmpMapTextureSize     = 1
    , tmpTilesetTextureSize = 1
    , tmpTileSize           = 1
    }

instance Storable TileMapParams where
  alignment ~_ = 4

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

  poke ptr TileMapParams{..} = do
    pokeByteOff ptr  0 tmpTextureIds
    pokeByteOff ptr 16 tmpViewOffset
    pokeByteOff ptr 24 tmpViewportSize
    pokeByteOff ptr 32 tmpMapTextureSize
    pokeByteOff ptr 40 tmpTilesetTextureSize
    pokeByteOff ptr 48 tmpTileSize

  peek ptr = do
    tmpTextureIds         <- peekByteOff ptr  0
    tmpViewOffset         <- peekByteOff ptr 16
    tmpViewportSize       <- peekByteOff ptr 24
    tmpMapTextureSize     <- peekByteOff ptr 32
    tmpTilesetTextureSize <- peekByteOff ptr 48
    tmpTileSize           <- peekByteOff ptr 56
    pure TileMapParams{..}

vkInstanceTileMap :: [Vk.Format]
vkInstanceTileMap =
  [ Vk.FORMAT_R32G32B32A32_SINT -- tmpTextureIds         :: IVec4
  , Vk.FORMAT_R32G32_SFLOAT     -- tmpViewOffset         :: Vec2
  , Vk.FORMAT_R32G32_SFLOAT     -- tmpViewportSize       :: Vec2
  , Vk.FORMAT_R32G32_SFLOAT     -- tmpMapTextureSize     :: Vec2
  , Vk.FORMAT_R32G32_SFLOAT     -- tmpTilesetTextureSize :: Vec2
  , Vk.FORMAT_R32G32_SFLOAT     -- tmpTileSize           :: Vec2
  ]

allocateInstancesWith
  :: ( MonadResource m
     , MonadUnliftIO m
     )
  => (Vk.BufferUsageFlagBits -> Int -> Storable.Vector TileMapParams -> m (InstanceTileMap 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 tilemapParams instances
    transforms = Storable.fromList $ map transformMat4 instances
    numInstances = Storable.length textures

    create (UnliftIO ul) = ul do
      ibTileMap   <- 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 ibTileMap
      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 TileMapParams, Storable.Vector Transform)
  -> m (InstanceBuffers 'Buffer.Coherent 'Buffer.Coherent)
updateCoherentResize_ context InstanceBuffers{..} (textures, transforms) =
  InstanceBuffers
    <$> Buffer.updateCoherentResize_ context ibTileMap textures
    <*> Buffer.updateCoherentResize_ context ibTransform transforms