module Engine.Types where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Control.Monad.Trans.Resource qualified as ResourceT
import Graphics.UI.GLFW qualified as GLFW
import RIO.App (App, appEnv, appState)
import RIO.Lens (_1)
import UnliftIO.Resource (MonadResource, ReleaseKey)
import Vulkan.Core10 qualified as Vk
import Vulkan.Extensions.VK_KHR_surface qualified as Khr
import Vulkan.NamedType ((:::))
import Vulkan.Utils.QueueAssignment (QueueFamilyIndex(..))
import VulkanMemoryAllocator qualified as VMA

import Engine.Camera (ProjectionProcess)
import Engine.Setup.Window (Window)
import Engine.Types.RefCounted (RefCounted)
import Engine.Vulkan.Swapchain (SwapchainResources(..))
import Engine.Vulkan.Types (HasVulkan(..))
import Engine.Vulkan.Types qualified as Vulkan

-- * App globals

-- | A bunch of global, unchanging state we cart around
data GlobalHandles = GlobalHandles
  { GlobalHandles -> Window
ghWindow             :: GLFW.Window
  , GlobalHandles -> SurfaceKHR
ghSurface            :: Khr.SurfaceKHR
  , GlobalHandles -> Instance
ghInstance           :: Vk.Instance
  , GlobalHandles -> PhysicalDevice
ghPhysicalDevice     :: Vk.PhysicalDevice
  , GlobalHandles -> PhysicalDeviceInfo
ghPhysicalDeviceInfo :: Vulkan.PhysicalDeviceInfo
  , GlobalHandles -> Device
ghDevice             :: Vk.Device
  , GlobalHandles -> Allocator
ghAllocator          :: VMA.Allocator
  , GlobalHandles -> Queues (QueueFamilyIndex, Queue)
ghQueues             :: Vulkan.Queues (QueueFamilyIndex, Vk.Queue)
  , GlobalHandles -> ProjectionProcess
ghScreenP            :: ProjectionProcess
  , GlobalHandles -> StageSwitchVar
ghStageSwitch        :: StageSwitchVar
  }

getScreenP :: MonadReader (App GlobalHandles s) m => m ProjectionProcess
getScreenP :: m ProjectionProcess
getScreenP = (App GlobalHandles s -> ProjectionProcess) -> m ProjectionProcess
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((App GlobalHandles s -> ProjectionProcess) -> m ProjectionProcess)
-> (App GlobalHandles s -> ProjectionProcess)
-> m ProjectionProcess
forall a b. (a -> b) -> a -> b
$ GlobalHandles -> ProjectionProcess
ghScreenP (GlobalHandles -> ProjectionProcess)
-> (App GlobalHandles s -> GlobalHandles)
-> App GlobalHandles s
-> ProjectionProcess
forall b c a. (b -> c) -> (a -> b) -> a -> c
. App GlobalHandles s -> GlobalHandles
forall env st. App env st -> env
appEnv

instance HasVulkan GlobalHandles where
  getInstance :: GlobalHandles -> Instance
getInstance           = GlobalHandles -> Instance
ghInstance
  getQueues :: GlobalHandles -> Queues (QueueFamilyIndex, Queue)
getQueues             = GlobalHandles -> Queues (QueueFamilyIndex, Queue)
ghQueues
  getPhysicalDevice :: GlobalHandles -> PhysicalDevice
getPhysicalDevice     = GlobalHandles -> PhysicalDevice
ghPhysicalDevice
  getPhysicalDeviceInfo :: GlobalHandles -> PhysicalDeviceInfo
getPhysicalDeviceInfo = GlobalHandles -> PhysicalDeviceInfo
ghPhysicalDeviceInfo
  getDevice :: GlobalHandles -> Device
getDevice             = GlobalHandles -> Device
ghDevice
  getAllocator :: GlobalHandles -> Allocator
getAllocator          = GlobalHandles -> Allocator
ghAllocator

-- * Stage stack

type StageStack = [StackStage]

data NextStage
  = Finish
  | Replace     StackStage
  | PushRestart StackStage
  | PushFreeze  StackStage

data StackStage where
  StackStage
    :: forall rp p rr st
    .  Vulkan.RenderPass rp
    => Stage rp p rr st
    -> StackStage
  StackStageContinue
    :: forall rp p rr st
    .  Vulkan.RenderPass rp
    => ReleaseKey
    -> st
    -> Stage rp p rr st
    -> StackStage

type StageSwitchVar = TMVar StageSwitch

data StageSwitch
  = StageSwitchPending NextStage
  | StageSwitchHandled

-- * Stage on a stack

type StageRIO st = RIO (App GlobalHandles st)

type StageSetupRIO = RIO (App GlobalHandles (Maybe SwapchainResources))

type StageFrameRIO rp p rr st = RIO (App GlobalHandles st, Frame rp p rr)

instance HasStateRef st (App GlobalHandles st, Frame rp p rr) where
  stateRefL :: (SomeRef st -> f (SomeRef st))
-> (App GlobalHandles st, Frame rp p rr)
-> f (App GlobalHandles st, Frame rp p rr)
stateRefL =
    ((App GlobalHandles st, Frame rp p rr) -> SomeRef st)
-> ((App GlobalHandles st, Frame rp p rr)
    -> SomeRef st -> (App GlobalHandles st, Frame rp p rr))
-> Lens' (App GlobalHandles st, Frame rp p rr) (SomeRef st)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens
      (App GlobalHandles st -> SomeRef st
forall env st. App env st -> SomeRef st
appState (App GlobalHandles st -> SomeRef st)
-> ((App GlobalHandles st, Frame rp p rr) -> App GlobalHandles st)
-> (App GlobalHandles st, Frame rp p rr)
-> SomeRef st
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (App GlobalHandles st, Frame rp p rr) -> App GlobalHandles st
forall a b. (a, b) -> a
fst)
      (\(App GlobalHandles st
app, Frame rp p rr
frame) SomeRef st
st' ->
          ( App GlobalHandles st
app
              { appState :: SomeRef st
appState = SomeRef st
st'
              }
          , Frame rp p rr
frame
          )
      )

data Stage rp p rr st = forall a . Stage
  { Stage rp p rr st -> Text
sTitle :: Text

  , Stage rp p rr st
-> SwapchainResources -> ResourceT (StageRIO st) rp
sAllocateRP :: SwapchainResources -> ResourceT (StageRIO st) rp
  , Stage rp p rr st
-> SwapchainResources -> rp -> ResourceT (StageRIO st) p
sAllocateP  :: SwapchainResources -> rp -> ResourceT (StageRIO st) p
  , Stage rp p rr st
-> StageRIO (Maybe SwapchainResources) (ReleaseKey, st)
sInitialRS  :: StageRIO (Maybe SwapchainResources) (ReleaseKey, st)
  , Stage rp p rr st
-> Queues CommandPool -> rp -> p -> ResourceT (StageRIO st) rr
sInitialRR  :: Vulkan.Queues Vk.CommandPool -> rp -> p -> ResourceT (StageRIO st) rr

  , ()
sBeforeLoop       :: StageRIO st a
  , Stage rp p rr st -> st -> rr -> StageFrameRIO rp p rr st ()
sUpdateBuffers    :: st -> rr -> StageFrameRIO rp p rr st ()
  , Stage rp p rr st
-> CommandBuffer
-> rr
-> ("image index" ::: Word32)
-> StageFrameRIO rp p rr st ()
sRecordCommands   :: Vk.CommandBuffer -> rr -> "image index" ::: Word32 -> StageFrameRIO rp p rr st ()
  , ()
sAfterLoop        :: a -> StageRIO st ()
  }

-- * Frame loop inside a stage

-- | All the information required to render a single frame
data Frame renderpass pipelines resources = Frame
  { Frame renderpass pipelines resources -> Word64
fIndex                       :: Word64 -- ^ Which number frame is this
  , Frame renderpass pipelines resources -> Window
fWindow                      :: Window
  , Frame renderpass pipelines resources -> SurfaceKHR
fSurface                     :: Khr.SurfaceKHR

  , Frame renderpass pipelines resources -> SwapchainResources
fSwapchainResources          :: SwapchainResources
  , Frame renderpass pipelines resources -> renderpass
fRenderpass                  :: renderpass
  , Frame renderpass pipelines resources -> pipelines
fPipelines                   :: pipelines
  , Frame renderpass pipelines resources -> Semaphore
fRenderFinishedHostSemaphore :: Vk.Semaphore
    {- ^
      A timeline semaphore which increments to fIndex when this frame
      is done, the host can wait on this semaphore.
    -}

  , Frame renderpass pipelines resources -> (RefCounted, InternalState)
fStageResources              :: (RefCounted, ResourceT.InternalState)
    -- ^ Swapchain-derived resources with a life time of this Frame's stage.

  , Frame renderpass pipelines resources -> IORef [GPUWork]
fGPUWork                     :: IORef [GPUWork]
    {- ^
      Timeline semaphores and corresponding wait values, updates as the
      frame progresses.
    -}

  , Frame renderpass pipelines resources -> (ReleaseKey, InternalState)
fResources                   :: (ReleaseKey, ResourceT.InternalState)
    {- ^
      The 'InternalState' for tracking frame-local resources along with the
      key to release it in the global scope. This will be released when the
      frame is done with GPU work.
    -}

  , Frame renderpass pipelines resources -> RecycledResources resources
fRecycledResources           :: RecycledResources resources
    {- ^
      Resources which can be used for this frame and are then passed on to a
      later frame.
    -}
  }

type GPUWork =
  ( "host semaphore" ::: Vk.Semaphore
  , "frame index" ::: Word64
  )

-- | These are resources which are reused by a later frame when the current
-- frame is retired
data RecycledResources a = RecycledResources
  { RecycledResources a -> Semaphore
rrImageAvailableSemaphore :: Vk.Semaphore
    -- ^ A binary semaphore passed to 'acquireNextImageKHR'
  , RecycledResources a -> Semaphore
rrRenderFinishedSemaphore :: Vk.Semaphore
    -- ^ A binary semaphore to synchronize rendering and presenting

  , RecycledResources a -> Queues CommandPool
rrQueues                  :: Vulkan.Queues Vk.CommandPool
    {- ^
      Pool for this frame's commands for each of the queue families.
      (might want more than one of these for multithreaded recording)
    -}

  , RecycledResources a -> a
rrData                    :: a
  }

instance HasLogFunc env => HasLogFunc (env, Frame rp p rr) where
  logFuncL :: (LogFunc -> f LogFunc)
-> (env, Frame rp p rr) -> f (env, Frame rp p rr)
logFuncL = (env -> f env) -> (env, Frame rp p rr) -> f (env, Frame rp p rr)
forall s t a b. Field1 s t a b => Lens s t a b
_1 ((env -> f env) -> (env, Frame rp p rr) -> f (env, Frame rp p rr))
-> ((LogFunc -> f LogFunc) -> env -> f env)
-> (LogFunc -> f LogFunc)
-> (env, Frame rp p rr)
-> f (env, Frame rp p rr)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LogFunc -> f LogFunc) -> env -> f env
forall env. HasLogFunc env => Lens' env LogFunc
logFuncL

instance MonadResource (RIO (env, Frame rp p rr)) where
  {-# INLINE liftResourceT #-}
  liftResourceT :: ResourceT IO a -> RIO (env, Frame rp p rr) a
liftResourceT ResourceT IO a
rt =
    ((env, Frame rp p rr) -> InternalState)
-> RIO (env, Frame rp p rr) InternalState
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((ReleaseKey, InternalState) -> InternalState
forall a b. (a, b) -> b
snd ((ReleaseKey, InternalState) -> InternalState)
-> ((env, Frame rp p rr) -> (ReleaseKey, InternalState))
-> (env, Frame rp p rr)
-> InternalState
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Frame rp p rr -> (ReleaseKey, InternalState)
forall renderpass pipelines resources.
Frame renderpass pipelines resources -> (ReleaseKey, InternalState)
fResources (Frame rp p rr -> (ReleaseKey, InternalState))
-> ((env, Frame rp p rr) -> Frame rp p rr)
-> (env, Frame rp p rr)
-> (ReleaseKey, InternalState)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (env, Frame rp p rr) -> Frame rp p rr
forall a b. (a, b) -> b
snd) RIO (env, Frame rp p rr) InternalState
-> (InternalState -> RIO (env, Frame rp p rr) a)
-> RIO (env, Frame rp p rr) a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
      IO a -> RIO (env, Frame rp p rr) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> RIO (env, Frame rp p rr) a)
-> (InternalState -> IO a)
-> InternalState
-> RIO (env, Frame rp p rr) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ResourceT IO a -> InternalState -> IO a
forall (m :: * -> *) a. ResourceT m a -> InternalState -> m a
ResourceT.runInternalState ResourceT IO a
rt