{-# OPTIONS_GHC -fplugin Foreign.Storable.Generic.Plugin #-}

{-# LANGUAGE OverloadedLists #-}

module Render.Unlit.Line2d.Model
  ( Segment

  , Vertex
  , createVertices
  , verticesRoundRound
  , Points

  , point
  , InstanceAttrs(..)

  , Buffer
  , Observer
  , Buffer.observeCoherentResize_

  , Batches(..)
  , BatchObserver
  , newBatchObserver
  , observeCoherentBatches
  ) where

import RIO

import Control.Monad.Trans.Resource qualified as Resource
import Data.Type.Equality (type (~))
import Data.Vector.Generic qualified as Generic
import Foreign.Storable.Generic (GStorable)
import Geomancy (Vec2, Vec4, vec3)
import Geomancy.Vec3 qualified as Vec3
import RIO.Vector qualified as Vector
import RIO.Vector.Storable qualified as Storable
import Vulkan.Core10 qualified as Vk
import Vulkan.NamedType ((:::))

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

type Segment = Buffer.Allocated 'Buffer.Staged Vec3.Packed
type Vertex = Model.Vertex3d ()

createVertices
  :: MonadVulkan env m
  => Maybe Text
  -> Queues Vk.CommandPool
  -> Float
  -> m Segment
createVertices label pools resolution = do
  Buffer.createStaged
    label
    pools
    Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT
    (Storable.length vertices)
    vertices
  where
    vertices = verticesRoundRound resolution

{- |
  Generate mesh for the round joints / round caps special case.

  With a bit of vertex shader code it allows drawing a batch of
  smooth lines in one call.
-}
verticesRoundRound :: Float -> Storable.Vector Vec3.Packed
verticesRoundRound resolution =
  Storable.fromList . map Vec3.Packed $
    mconcat
      [ segment
      , leftSemi
      , rightSemi
      ]
  where
    segment =
      [ vec3 0 (-0.5) 0
      , vec3 0 (-0.5) 1
      , vec3 0   0.5  1
      , vec3 0 (-0.5) 0
      , vec3 0   0.5  1
      , vec3 0   0.5  0
      ]

    leftSemi = do
      step <- [0..resolution]
      let
        theta0 = pi / 2 + (step + 0) * pi / resolution
        theta1 = pi / 2 + (step + 1) * pi / resolution
        a = vec3 0 0 0
        b = vec3 (cos theta0) (sin theta0) 0 * 0.5
        c = vec3 (cos theta1) (sin theta1) 0 * 0.5
      [ a, b, c ]

    rightSemi = do
      step <- [0..resolution]
      let
        theta0 = 3 * pi / 2 + (step + 0) * pi / resolution
        theta1 = 3 * pi / 2 + (step + 1) * pi / resolution
        a = vec3 0 0 0
        b = vec3 (cos theta0) (sin theta0) 2 * 0.5
        c = vec3 (cos theta1) (sin theta1) 2 * 0.5
      [ a, b, c ]

type Points = Storable.Vector InstanceAttrs

point :: Float -> Vec4 -> Vec2 -> InstanceAttrs
point width color position = InstanceAttrs
  { color    = color
  , position = Vec3.Packed (Vec3.fromVec2 position 0)
  , width    = width
  }

data InstanceAttrs = InstanceAttrs
  { color    :: Vec4

  , position :: Vec3.Packed
  , width    :: Float
  } deriving (Eq, Show, Generic)

-- XXX: okay, the layout matches
instance GStorable InstanceAttrs

instance HasVkFormat InstanceAttrs where
  getVkFormat =
    [ Vk.FORMAT_R32G32B32A32_SFLOAT -- color
    , Vk.FORMAT_R32G32B32A32_SFLOAT -- position+width
    ]

instance HasVertexInputBindings InstanceAttrs where
  vertexInputBindings =
    -- XXX: instance buffer is bound 2 times with a shift
    replicate 2 $ instanceFormat @InstanceAttrs

type Buffer s = Buffer.Allocated s InstanceAttrs

type Observer = Buffer.ObserverCoherent InstanceAttrs
type BatchObserver = Worker.ObserverIO (Buffer 'Buffer.Coherent, Ranges)
type Ranges = [(Word32, Word32)]

newtype Batches v a = Batches [v a]
  deriving (Eq, Show, Ord, Functor, Foldable, Traversable, Semigroup, Monoid)

newBatchObserver
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     )
  => "initial size" ::: Int
  -> m BatchObserver
newBatchObserver initialSize = do
  initialBuffer <- Buffer.createCoherent
    (Just "Line2D")
    Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT
    initialSize
    mempty
  observer <- Worker.newObserverIO (initialBuffer, [])

  context <- ask
  void $! Resource.register do
    (currentBuffer, _ranges) <- Worker.readObservedIO observer
    Buffer.destroy context currentBuffer

  pure (observer :: BatchObserver)

observeCoherentBatches
  :: ( Worker.GetOutput output ~ Batches Storable.Vector InstanceAttrs
     , Worker.HasOutput output
     , MonadVulkan env m
     )
  => output
  -> BatchObserver
  -> m ()
observeCoherentBatches source observer =
  Worker.observeIO_ source observer \(buf, _ranges) (Batches batches) -> do
    let segments = filterSegments batches
    buf' <- Buffer.updateCoherentResize_ buf $
      Storable.concat segments
    pure (buf', toRanges segments)

filterSegments
  :: Generic.Vector v a
  => [v a]
  -> [v a]
filterSegments = filter \v ->
  Vector.length v >= 2

toRanges
  :: ( Generic.Vector v a
     , Integral i
     )
  => [v a]
  -> [(i, i)]
toRanges = collect . filter (not . Vector.null)
  where
    collect v = Vector.toList @Vector $ Vector.zip offsets sizes
      where
        sizes = Vector.fromList $ map (fromIntegral . Vector.length) v
        offsets = Vector.scanl' (+) 0 sizes