module Render.Draw
  ( triangle_
  , triangles_
  , quads

  , indexed
  , indexedRanges
  , indexedParts

  , indexedPos
  , indexedPosRanges

  , unsafeIndexedRanges
  , unsafeIndexedParts
  ) where

import RIO

import Data.Vector qualified as Vector
import Vulkan.Core10 qualified as Vk

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

-- * Primitives

-- | Single triangle, binding nothing.
triangle_ :: MonadUnliftIO m => Vk.CommandBuffer -> Bound dsl () () m ()
triangle_ cb = triangles_ cb 1

-- | Multiple shader-driven triangles without bindings.
triangles_ :: MonadUnliftIO m => Vk.CommandBuffer -> Word32 -> Bound dsl () () m ()
triangles_ cb num =
  Bound $ Vk.cmdDraw cb 3 num 0 0

-- | Instanced quads.
quads
  :: MonadUnliftIO m
  => Vk.CommandBuffer
  -> Buffer.Allocated stage instances
  -> Bound dsl () instances m ()
quads cb instances = Bound do
  Vk.cmdBindVertexBuffers cb 0 (pure $ Buffer.aBuffer instances) (pure 0)
  Vk.cmdDraw cb 6 (Buffer.aUsed instances) 0 0

-- * Indexed models

-- ** Positions + attributes

-- | Draw whole-model instances.
indexed
  :: (MonadUnliftIO m, Model.HasVertexBuffers instances)
  => Vk.CommandBuffer
  -> Model.Indexed storage pos attrs
  -> instances
  -> Bound dsl (Model.Vertex pos attrs) (Model.VertexBuffersOf instances) m ()
indexed cmd model instances = indexedRanges cmd model instances [wholeIndexed]
  where
    wholeIndexed = Model.IndexRange
      { irFirstIndex = 0
      , irIndexCount = Buffer.aUsed (Model.iIndices model)
      }

{- |
  Draw subrange of each instance.

  E.g. chunks of the same material drawn in different places.
-}
indexedRanges
  :: (MonadUnliftIO m, Model.HasVertexBuffers instances)
  => Vk.CommandBuffer
  -> Model.Indexed storage pos attrs
  -> instances
  -> [Model.IndexRange]
  -> Bound dsl (Model.Vertex pos attrs) (Model.VertexBuffersOf instances) m ()
indexedRanges cmd model instances ranges = do
  checkedRanges <- traverse check ranges
  Bound $ unsafeIndexedRanges True cmd model instances checkedRanges
  where
    check ir@Model.IndexRange{..}
      | irFirstIndex > maxIndex =
          throwString "firstIndex is over the actual buffer size"
      | irFirstIndex + irIndexCount > maxIndex =
          throwString "firstIndex + indexCount is over the actual buffer size"
      | otherwise =
          pure ir

    maxIndex = Buffer.aUsed (Model.iIndices model)

{- |
  Draw ranges and instances zipped.

  E.g. range materials stored in instances.
-}
indexedParts
  :: (MonadUnliftIO m, Model.HasVertexBuffers instances, Foldable t)
  => Bool
  -> Vk.CommandBuffer
  -> Model.Indexed storage pos attrs
  -> instances
  -> Int
  -> t Model.IndexRange
  -> Bound dsl (Model.Vertex pos attrs) (Model.VertexBuffersOf instances) m ()
indexedParts drawAttrs cmd model instances startInstance parts =
  Bound $ unsafeIndexedParts drawAttrs cmd model instances startInstance parts

-- ** Position-only draws

-- | Draw whole-model instances, ignoring attributes.
indexedPos
  :: (MonadUnliftIO m, Model.HasVertexBuffers instances)
  => Vk.CommandBuffer
  -> Model.Indexed storage pos unusedAttrs
  -> instances
  -> Bound dsl (Model.Vertex pos ignoreAttrs) (Model.VertexBuffersOf instances) m ()
indexedPos cmd model instances = indexedPosRanges cmd model instances [wholeIndexed]
  where
    wholeIndexed = Model.IndexRange
      { irFirstIndex = 0
      , irIndexCount = Buffer.aUsed (Model.iIndices model)
      }

-- | Draw subrange of each instances, ignoring attributes.
indexedPosRanges
  :: (MonadUnliftIO m, Model.HasVertexBuffers instances)
  => Vk.CommandBuffer
  -> Model.Indexed storage pos unusedAttrs
  -> instances
  -> [Model.IndexRange]
  -> Bound dsl (Model.Vertex pos ignoreAttrs) (Model.VertexBuffersOf instances) m ()
indexedPosRanges cmd model instances ranges = do
  checkedRanges <- traverse check ranges
  Bound $ unsafeIndexedRanges False cmd model instances checkedRanges
  where
    check ir@Model.IndexRange{..}
      | irFirstIndex > maxIndex =
          throwString "firstIndex is over the actual buffer size"
      | irFirstIndex + irIndexCount > maxIndex =
          throwString "firstIndex + indexCount is over the actual buffer size"
      | otherwise =
          pure ir

    maxIndex = Buffer.aUsed (Model.iIndices model)

-- * Unchecked implementations

-- | Common unchecked part for pos/attrs
unsafeIndexedRanges
  :: (MonadUnliftIO io, Model.HasVertexBuffers instances, Foldable t)
  => Bool
  -> Vk.CommandBuffer
  -> Model.Indexed storage pos attrs
  -> instances
  -> t Model.IndexRange
  -> io ()
unsafeIndexedRanges drawAttrs cmd Model.Indexed{..} instances indexRanges =
  case toList indexRanges of
    [] ->
      pure ()

    _skip | instanceCount < 1 ->
      pure ()

    someRanges -> liftIO do
      Vk.cmdBindVertexBuffers cmd 0 vertexBuffers bufferOffsets
      Vk.cmdBindIndexBuffer cmd (Buffer.aBuffer iIndices) indexBufferOffset Vk.INDEX_TYPE_UINT32
      for_ someRanges \Model.IndexRange{..} ->
        Vk.cmdDrawIndexed cmd irIndexCount instanceCount irFirstIndex vertexOffset firstInstance

  where
    indexBufferOffset = 0
    vertexOffset = 0

    instanceCount = Model.getInstanceCount instances
    firstInstance = 0

    vertexBuffers = Vector.fromList $
      if drawAttrs then
        Buffer.aBuffer iPositions :
        Buffer.aBuffer iAttrs :
        Model.getVertexBuffers instances
      else
        Buffer.aBuffer iPositions :
        Model.getVertexBuffers instances

    bufferOffsets =
      Vector.replicate (Vector.length vertexBuffers) 0

-- | Instance/range zipped
unsafeIndexedParts
  :: (MonadUnliftIO io, Model.HasVertexBuffers instances, Foldable t)
  => Bool
  -> Vk.CommandBuffer
  -> Model.Indexed storage pos attrs
  -> instances
  -> Int
  -> t Model.IndexRange
  -> io ()
unsafeIndexedParts drawAttrs cmd Model.Indexed{..} instances startInstance parts =
  case drop startInstance (toList parts) of
    [] ->
      pure ()

    _skip | instanceCount < 1 ->
      pure ()

    someRanges -> liftIO do
      Vk.cmdBindVertexBuffers cmd 0 vertexBuffers bufferOffsets
      Vk.cmdBindIndexBuffer cmd (Buffer.aBuffer iIndices) indexBufferOffset Vk.INDEX_TYPE_UINT32
      for_ (zip [fromIntegral startInstance ..] someRanges) \(firstInstance, Model.IndexRange{..}) ->
        Vk.cmdDrawIndexed cmd irIndexCount instanceCount irFirstIndex vertexOffset firstInstance

  where
    indexBufferOffset = 0
    vertexOffset = 0

    instanceCount = 1 -- Model.getInstanceCount instances
    -- firstInstance = 0

    vertexBuffers = Vector.fromList $
      if drawAttrs then
        Buffer.aBuffer iPositions :
        Buffer.aBuffer iAttrs :
        Model.getVertexBuffers instances
      else
        Buffer.aBuffer iPositions :
        Model.getVertexBuffers instances

    bufferOffsets =
      Vector.replicate (Vector.length vertexBuffers) 0