{- | Experimental helpers for managing models with multiple instance buffers.

Only works with fully-coherent models and atomic stores.
Not particularly efficient: when any element is changed, everything gets fully updated.
-}

module Resource.Model.Observer
  ( newCoherent
  , observeCoherent

  , VertexBuffers(..)
  , genericCreateInitial
  , genericDestroyCurrent

  , UpdateCoherent(..)
  , genericUpdateCoherent
  ) where

import RIO
import GHC.Generics

import Control.Monad.Trans.Resource (ResourceT)
import Control.Monad.Trans.Resource qualified as Resource
import Data.Kind (Type)
import Data.Text qualified as Text
import Data.Vector.Storable qualified as Storable
import Vulkan.Core10 qualified as Vk

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

newCoherent
  :: ( VertexBuffers res
     , MonadVulkan env m
     )
  => Int
  -> Text
  -> ResourceT m (Worker.ObserverIO res)
newCoherent initialSize label = do
  initial <- createInitial initialSize label
  observer <- Worker.newObserverIO initial

  context <- ask
  void $! Resource.register do
    current <- Worker.readObservedIO observer
    destroyCurrent context current

  pure observer

observeCoherent
  :: ( MonadVulkan env m
    , Worker.HasOutput output
    , UpdateCoherent bufs (Worker.GetOutput output)
    )
  => output
  -> Worker.ObserverIO bufs
  -> m ()
observeCoherent process observer =
  void $ Worker.observeIO process observer updateCoherent

class VertexBuffers a where
  createInitial :: forall env m . MonadVulkan env m => Int -> Text -> ResourceT m a

  default createInitial
    :: forall env m
    .  ( Generic a
       , GVertexBuffers (Rep a)
       , MonadVulkan env m
       )
    => Int
    -> Text
    -> ResourceT m a
  createInitial = genericCreateInitial

  destroyCurrent :: forall env . HasVulkan env => env -> a -> IO ()
  default destroyCurrent
    :: forall env
    .  ( Generic a
       , GVertexBuffers (Rep a)
       , HasVulkan env
       )
    => env
    -> a
    -> IO ()
  destroyCurrent = genericDestroyCurrent

{- XXX: Terminal instance for generic observable requiring only MonadVulkan.

Initial values can be written post-hoc using `updateCoherent` on observed data.
The more problematic part is using hardcoded usage argument.
This is fine for the main use case of model instance buffers, but makes it unsuitable for other uses.
-}
instance Storable a => VertexBuffers (Buffer.Allocated 'Buffer.Coherent a) where
  createInitial initialCapacity label =
    Buffer.createCoherent
      (Just label)
      Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT
      initialCapacity
      mempty

  destroyCurrent ctx buf = Buffer.destroy ctx buf

genericCreateInitial
  :: ( Generic a
     , GVertexBuffers (Rep a)
     , MonadVulkan env m
     )
  => Int
  -> Text
  -> ResourceT m a
genericCreateInitial size label =
  fmap GHC.Generics.to (gcreate size label)

genericDestroyCurrent
  :: ( Generic a
     , GVertexBuffers (Rep a)
     , HasVulkan env
     )
  => env
  -> a
  -> IO ()
genericDestroyCurrent env = gdestroy env . GHC.Generics.from

class GVertexBuffers (f :: Type -> Type) where
  gcreate :: forall env m a . MonadVulkan env m => Int -> Text -> ResourceT m (f a)
  gdestroy :: forall env a . HasVulkan env => env -> f a -> IO ()

instance GVertexBuffers f => GVertexBuffers (D1 c f) where
  gcreate size label = fmap M1 $ gcreate size label
  gdestroy env (M1 a) = gdestroy env a

instance GVertexBuffers f => GVertexBuffers (C1 c f) where
  gcreate size label = fmap M1 $ gcreate size label
  gdestroy env (M1 a) = gdestroy env a

instance (GVertexBuffers l, GVertexBuffers r) => GVertexBuffers (l :*: r) where
  gcreate size label = (:*:)
    <$> gcreate size label
    <*> gcreate size label
  gdestroy env (l :*: r) =
    gdestroy env r <>
    gdestroy env l

instance (GVertexBuffers f, Selector c) => GVertexBuffers (M1 S c f) where
  gcreate size label =
    fmap M1 $
      gcreate size $
        Text.intercalate "."
          [ label
          , fromString $ selName (undefined :: t c f a)
          ]
  gdestroy env (M1 a) = gdestroy env a

instance VertexBuffers a => GVertexBuffers (K1 r a) where
  gcreate size label = fmap K1 $ createInitial size label
  gdestroy env (K1 a) = destroyCurrent env a

-- TODO: class UpdateStaged bufs stores

class UpdateCoherent bufs stores where
  updateCoherent :: forall env m . MonadVulkan env m => bufs -> stores -> m bufs

  default updateCoherent
    :: forall env m
    .  ( Generic bufs
       , Generic stores
       , GUpdateCoherent (Rep bufs) (Rep stores)
       , MonadVulkan env m
       )
    => bufs
    -> stores
    -> m bufs
  updateCoherent = genericUpdateCoherent

genericUpdateCoherent
  :: ( Generic bufs
      , Generic stores
      , GUpdateCoherent (Rep bufs) (Rep stores)
      , MonadVulkan env m
      )
  => bufs
  -> stores
  -> m bufs
genericUpdateCoherent bufs stores =
  fmap GHC.Generics.to $
    gUpdateCoherent
      (GHC.Generics.from bufs)
      (GHC.Generics.from stores)

instance Storable a => UpdateCoherent (Buffer.Allocated 'Buffer.Coherent a) (Storable.Vector a) where
  updateCoherent = Buffer.updateCoherentResize_

-- TODO: barbies-style update
-- instance (ApplicativeB collection) => UpdateCoherent (collection (Buffer.Allocated 'Buffer.Coherent)) (collection Storable.Vector) where
--   update = bzipWith Buffer.updateCoherentResize_

class GUpdateCoherent (fb :: Type -> Type) (fs :: Type -> Type) where
  gUpdateCoherent :: forall env m b s . MonadVulkan env m => fb b -> fs s -> m (fb b)

instance GUpdateCoherent fb fs => GUpdateCoherent (M1 c cb fb) (M1 c cs fs) where
  gUpdateCoherent (M1 gb) (M1 gs) = fmap M1 $ gUpdateCoherent gb gs

instance (GUpdateCoherent fbl fsl, GUpdateCoherent fbr fsr) => GUpdateCoherent (fbl :*: fbr) (fsl :*: fsr) where
  gUpdateCoherent (bl :*: br) (sl :*: sr) = (:*:)
    <$> gUpdateCoherent bl sl
    <*> gUpdateCoherent br sr

instance UpdateCoherent ba sa => GUpdateCoherent (K1 br ba) (K1 sr sa) where
  gUpdateCoherent (K1 gb) (K1 gs) = fmap K1 $ updateCoherent gb gs