{-# LANGUAGE AllowAmbiguousTypes #-}

module Resource.Model where

import RIO
import GHC.Generics

import Codec.Serialise qualified as CBOR
import Data.Kind (Constraint, Type)
import Data.List qualified as List
import Data.Vector.Storable qualified as Storable
import Foreign (Storable(..))
import Geomancy (Vec2)
import Geomancy.Vec3 qualified as Vec3
import UnliftIO.Resource (MonadResource)
import Vulkan.Core10 qualified as Vk

import Engine.Vulkan.Pipeline.Graphics (HasVertexInputBindings(..))
import Engine.Vulkan.Format (HasVkFormat(..))
import Engine.Vulkan.Pipeline.Graphics qualified as Graphics
import Engine.Vulkan.Types (MonadVulkan, Queues(..))
import Resource.Buffer qualified as Buffer

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

type Vertex2d attrs = Vertex Vec2 attrs
type Vertex3d attrs = Vertex Vec3.Packed attrs

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
    }

instance
  ( HasVkFormat pos
  , HasVkFormat attrs
  ) =>
  Graphics.HasVertexInputBindings (Vertex pos attrs) where
    vertexInputBindings =
      [ Graphics.vertexFormat @pos
      , Graphics.vertexFormat @attrs
      ]

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

  default getVertexBuffers
    :: ( Generic a
       , GHasVertexBuffers (Rep a)
       )
    => a
    -> [Vk.Buffer]
  getVertexBuffers = genericGetVertexBuffers

  default getInstanceCount
    :: ( Generic a
       , GHasVertexBuffers (Rep a)
       )
    => a
    -> Word32
  getInstanceCount = genericGetInstanceCount

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

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

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

instance forall (a :: Type) (store :: Buffer.Store) . 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

genericGetVertexBuffers
  :: ( Generic a
     , GHasVertexBuffers (Rep a)
     )
  => a
  -> [Vk.Buffer]
genericGetVertexBuffers = gVertexBuffers . GHC.Generics.from

genericGetInstanceCount
  :: ( Generic a
     , GHasVertexBuffers (Rep a)
     )
  => a
  -> Word32
genericGetInstanceCount = gInstanceCount . GHC.Generics.from

type GHasVertexBuffers :: (Type -> Type) -> Constraint
class GHasVertexBuffers f where
  gVertexBuffers  :: forall a . f a -> [Vk.Buffer]
  gInstanceCount :: forall a . f a -> Word32

instance GHasVertexBuffers f => GHasVertexBuffers (M1 c cb f) where
  gVertexBuffers (M1 f) = gVertexBuffers f
  gInstanceCount (M1 f) = gInstanceCount f

instance (GHasVertexBuffers l, GHasVertexBuffers r) => GHasVertexBuffers (l :*: r) where
  gVertexBuffers (l :*: r) = gVertexBuffers l <> gVertexBuffers r
  gInstanceCount (l :*: r) = min (gInstanceCount l) (gInstanceCount r)

instance HasVertexBuffers a => GHasVertexBuffers (K1 r a) where
  gVertexBuffers (K1 a) = getVertexBuffers a
  gInstanceCount (K1 a) = getInstanceCount a

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

instance CBOR.Serialise IndexRange

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
  :: ( MonadVulkan env m
     , Storable pos
     , Storable attrs
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> [Vertex pos attrs]
  -> Maybe [Word32]
  -> m (Indexed 'Buffer.Staged pos attrs)
createStagedL label pool vertices mindices = createStaged label 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
  :: ( MonadVulkan env m
     , Storable pos
     , Storable attrs
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> Storable.Vector pos
  -> Storable.Vector attrs
  -> Storable.Vector Word32
  -> m (Indexed 'Buffer.Staged pos attrs)
createStaged label pool pv av iv = do
  positions <- Buffer.createStaged (label <&> \m -> mappend m ".positions") pool Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT 0 pv
  attrs     <- Buffer.createStaged (label <&> \m -> mappend m ".attrs") pool Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT 0 av
  indices   <- Buffer.createStaged (label <&> \m -> mappend m ".indices") pool Vk.BUFFER_USAGE_INDEX_BUFFER_BIT  0 iv
  pure Indexed
    { iLabel     = label
    , iPositions = positions
    , iAttrs     = attrs
    , iIndices   = indices
    }

createCoherentEmpty
  :: ( MonadVulkan env m
     , Storable pos
     , Storable attrs
     )
  => Maybe Text
  -> Int
  -> m (Indexed 'Buffer.Coherent pos attrs)
createCoherentEmpty label initialSize = Indexed label
  <$> Buffer.createCoherent (label <&> \m -> mappend m ".positions") Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT initialSize mempty
  <*> Buffer.createCoherent (label <&> \m -> mappend m ".attrs") Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT initialSize mempty
  <*> Buffer.createCoherent (label <&> \m -> mappend m ".indices") Vk.BUFFER_USAGE_INDEX_BUFFER_BIT  initialSize mempty

registerIndexed_
  :: ( MonadVulkan env m
     , MonadResource m
     )
  => Indexed storage pos attrs
  -> m ()
registerIndexed_ Indexed{..} = do
  void $! Buffer.register iPositions
  void $! Buffer.register iAttrs
  void $! Buffer.register iIndices

destroyIndexed
  :: MonadVulkan env m
  => Indexed storage pos attrs
  -> m ()
destroyIndexed Indexed{..} = do
  ctx <- ask
  Buffer.destroy ctx iPositions
  Buffer.destroy ctx iAttrs
  Buffer.destroy ctx iIndices

updateCoherent
  :: ( MonadVulkan env m
     , Storable pos
     , Storable attrs
     )
  => [Vertex pos attrs]
  -> Indexed 'Buffer.Coherent pos attrs
  -> m (Indexed 'Buffer.Coherent pos attrs)
updateCoherent vertices old = do
  Indexed{..} <- pick
  Indexed iLabel
    <$> Buffer.updateCoherent pv iPositions
    <*> Buffer.updateCoherent av iAttrs
    <*> Buffer.updateCoherent iv iIndices
  where
    pick =
      if oldSize > newSize then
        pure old
      else do
        destroyIndexed old
        createCoherentEmpty (iLabel old) (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)