module Engine.Stage.Component where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Data.Semigroup (Semigroup(..))
import UnliftIO.Resource (ReleaseKey, allocate)
import Vulkan.Core10 qualified as Vk
import Vulkan.NamedType ((:::))

import Engine.Types (Stage(..), StageFrameRIO, StageRIO)
import Engine.Vulkan.Swapchain (SwapchainResources(..))
import Engine.Vulkan.Types qualified as Vulkan
import Resource.Region qualified as Region

assemble
  :: Foldable t
  => Text
  -> Rendering rp p st
  -> Resources rp p st rr
  -> t (Scene rp p st rr)
  -> Stage rp p rr st
assemble title Rendering{..} Resources{..} (fold -> Scene{..}) = Stage
  { sTitle = title

  , sAllocateRP = rAllocateRP
  , sAllocateP  = rAllocateP

  , sInitialRS  = rInitialRS
  , sInitialRR  = rInitialRR

  , sBeforeLoop       = Region.exec scBeforeLoop
  , sUpdateBuffers    = scUpdateBuffers
  , sRecordCommands   = scRecordCommands
  , sAfterLoop        = Region.release
  }

data Rendering rp p st = Rendering
  { rAllocateRP :: SwapchainResources -> ResourceT (StageRIO st) rp
  , rAllocateP  :: SwapchainResources -> rp -> ResourceT (StageRIO st) p
  }

data  NoRenderPass = NoRenderPass

instance Vulkan.RenderPass NoRenderPass where
  updateRenderpass _context = pure
  refcountRenderpass _rp = pure ()

data NoPipelines = NoPipelines
type  NoRendering = Rendering NoRenderPass NoPipelines

noRendering :: NoRendering st
noRendering = Rendering
  { rAllocateRP = \_swapchain ->
      pure NoRenderPass
  , rAllocateP = \_swapchain NoRenderPass ->
      pure NoPipelines
  }

data Resources rp p st rr = Resources
  { rInitialRS :: StageRIO (Maybe SwapchainResources) (ReleaseKey, st)
  , rInitialRR :: Vulkan.Queues Vk.CommandPool -> rp -> p -> ResourceT (StageRIO st) rr
  }

type NoResources rp p = Resources rp p NoRunState NoFrameResources

data NoRunState = NoRunState
data NoFrameResources = NoFrameResources

noResources :: NoResources rp p
noResources = Resources
  { rInitialRS =
      allocate
        (pure NoRunState)
        (\NoRunState -> pure ())
  , rInitialRR = \_pool _rp _p ->
      pure NoFrameResources
  }

data Scene rp p st rr = Scene
  { scBeforeLoop       :: ResourceT (StageRIO st) ()
  , scUpdateBuffers    :: st -> rr -> StageFrameRIO rp p rr st ()
  , scRecordCommands   :: Vk.CommandBuffer -> rr -> "image index" ::: Word32 -> StageFrameRIO rp p rr st ()
  }

instance Semigroup (Scene rp p st rr) where
  a <> b = Scene
    { scBeforeLoop = do
        scBeforeLoop a
        scBeforeLoop b
    , scUpdateBuffers = \st rr -> do
        scUpdateBuffers a st rr
        scUpdateBuffers b st rr
    , scRecordCommands = \cb rr ii -> do
        scRecordCommands a cb rr ii
        scRecordCommands b cb rr ii
    }

  sconcat scenes = Scene
    { scBeforeLoop = do
        traverse_ scBeforeLoop scenes
    , scUpdateBuffers = \st rr ->
        for_ scenes \Scene{scUpdateBuffers} ->
          scUpdateBuffers st rr
    , scRecordCommands = \cb rr ii ->
        for_ scenes \Scene{scRecordCommands} ->
          scRecordCommands cb rr ii
    }

instance Monoid (Scene rp p st rr) where
  mempty = Scene
    { scBeforeLoop = pure ()
    , scUpdateBuffers = mempty
    , scRecordCommands = mempty
    }