module Resource.Model where

import RIO

import Data.List qualified as List
import Data.Vector.Storable qualified as Storable
import Foreign (Storable(..))
import Vulkan.Core10 qualified as Vk

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

data Indexed storage pos attrs = Indexed
  { iPositions :: Buffer.Allocated storage pos
  , iAttrs     :: Buffer.Allocated storage attrs
  , iIndices   :: Buffer.Allocated storage Word32
  } deriving (Show)

data Vertex pos attrs = Vertex
  { vPosition :: pos
  , vAttrs    :: attrs
  } deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

{-# INLINEABLE vertexAttrs #-}
vertexAttrs :: (pos -> a -> b) -> [Vertex pos a] -> [Vertex pos b]
vertexAttrs inject vertices = do
  v@Vertex{..} <- vertices
  pure v
    { vAttrs = inject vPosition vAttrs
    }

{-# INLINEABLE vertexAttrsPos #-}
vertexAttrsPos :: (pos -> a) -> [pos] -> [Vertex pos a]
vertexAttrsPos inject positions = do
  pos <- positions
  pure Vertex
    { vPosition = pos
    , vAttrs    = inject pos
    }

class HasVertexBuffers a where
  type VertexBuffersOf a
  getVertexBuffers :: a -> [Vk.Buffer]
  getInstanceCount :: a -> Word32

instance HasVertexBuffers () where
  type VertexBuffersOf () = ()

  {-# INLINE getVertexBuffers #-}
  getVertexBuffers () = []

  {-# INLINE getInstanceCount #-}
  getInstanceCount () = 1

instance HasVertexBuffers (Buffer.Allocated store a) where
  type VertexBuffersOf (Buffer.Allocated store a) = a

  {-# INLINE getVertexBuffers #-}
  getVertexBuffers Buffer.Allocated{aBuffer} = [aBuffer]

  {-# INLINE getInstanceCount #-}
  getInstanceCount = Buffer.aUsed

data IndexRange = IndexRange
  { irFirstIndex :: Word32
  , irIndexCount :: Word32
  }
  deriving (Eq, Ord, Show)

instance Storable IndexRange where
  alignment ~_ = 4

  sizeOf ~_ = 8

  peek ptr = do
    irFirstIndex <- peekByteOff ptr 0
    irIndexCount <- peekByteOff ptr 4
    pure IndexRange{..}

  poke ptr IndexRange{..} = do
    pokeByteOff ptr 0 irFirstIndex
    pokeByteOff ptr 4 irIndexCount

createStagedL
  :: (HasVulkan context, Storable pos, Storable attrs, MonadUnliftIO io)
  => context
  -> Queues Vk.CommandPool
  -> [Vertex pos attrs]
  -> Maybe [Word32]
  -> io (Indexed 'Buffer.Staged pos attrs)
createStagedL context pool vertices mindices = createStaged context pool pv av iv
  where
    pv = Storable.fromList ps
    av = Storable.fromList as

    iv = case mindices of
      Just is ->
        Storable.fromList is
      Nothing ->
        -- TODO: add vertex deduplication
        Storable.generate (Storable.length pv) fromIntegral

    (ps, as) = List.unzip do
      Vertex{..} <- vertices
      pure (vPosition, vAttrs)

createStaged
  :: (HasVulkan context, Storable pos, Storable attrs, MonadUnliftIO io)
  => context
  -> Queues Vk.CommandPool
  -> Storable.Vector pos
  -> Storable.Vector attrs
  -> Storable.Vector Word32
  -> io (Indexed 'Buffer.Staged pos attrs)
createStaged context pool pv av iv = do
  positions <- Buffer.createStaged context pool Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT 0 pv
  attrs     <- Buffer.createStaged context pool Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT 0 av
  indices   <- Buffer.createStaged context pool Vk.BUFFER_USAGE_INDEX_BUFFER_BIT  0 iv
  pure Indexed
    { iPositions = positions
    , iAttrs     = attrs
    , iIndices   = indices
    }

createCoherentEmpty
  :: (HasVulkan context, Storable pos, Storable attrs, MonadUnliftIO io)
  => context
  -> Int
  -> io (Indexed 'Buffer.Coherent pos attrs)
createCoherentEmpty ctx initialSize = Indexed
  <$> Buffer.createCoherent ctx Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT initialSize mempty
  <*> Buffer.createCoherent ctx Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT initialSize mempty
  <*> Buffer.createCoherent ctx Vk.BUFFER_USAGE_INDEX_BUFFER_BIT  initialSize mempty

destroyIndexed
  :: (HasVulkan context, MonadUnliftIO io)
  => context
  -> Indexed storage pos attrs
  -> io ()
destroyIndexed ctx Indexed{..} = do
  Buffer.destroyAll ctx [iPositions]
  Buffer.destroyAll ctx [iAttrs]
  Buffer.destroyAll ctx [iIndices]

updateCoherent
  :: (HasVulkan context, Storable pos, Storable attrs, MonadUnliftIO io)
  => context
  -> [Vertex pos attrs]
  -> Indexed 'Buffer.Coherent pos attrs
  -> io (Indexed 'Buffer.Coherent pos attrs)
updateCoherent ctx vertices old = do
  Indexed{..} <- pick
  Indexed
    <$> Buffer.updateCoherent pv iPositions
    <*> Buffer.updateCoherent av iAttrs
    <*> Buffer.updateCoherent iv iIndices
  where
    pick =
      if oldSize > newSize then
        pure old
      else do
        destroyIndexed ctx old
        createCoherentEmpty ctx (max newSize $ oldSize * 2)

    oldSize = Buffer.aCapacity $ iIndices old
    newSize = Storable.length pv

    pv = Storable.fromList ps
    av = Storable.fromList as

    iv = Storable.generate newSize fromIntegral

    (ps, as) = List.unzip do
      Vertex{..} <- vertices
      pure (vPosition, vAttrs)