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.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
import Engine.Worker qualified as Worker
import Engine.Types.Options (Options)

-- * App globals

-- | A bunch of global, unchanging state we cart around
data GlobalHandles = GlobalHandles
  { GlobalHandles -> Options
ghOptions            :: Options
  , 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 -> Var Extent2D
ghScreenVar          :: Worker.Var Vk.Extent2D
  , GlobalHandles -> StageSwitchVar
ghStageSwitch        :: StageSwitchVar
  }

askScreenVar :: StageRIO env (Worker.Var Vk.Extent2D)
askScreenVar :: forall env. StageRIO env (Var Extent2D)
askScreenVar = (App GlobalHandles env -> Var Extent2D)
-> RIO (App GlobalHandles env) (Var Extent2D)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((App GlobalHandles env -> Var Extent2D)
 -> RIO (App GlobalHandles env) (Var Extent2D))
-> (App GlobalHandles env -> Var Extent2D)
-> RIO (App GlobalHandles env) (Var Extent2D)
forall a b. (a -> b) -> a -> b
$ GlobalHandles -> Var Extent2D
ghScreenVar (GlobalHandles -> Var Extent2D)
-> (App GlobalHandles env -> GlobalHandles)
-> App GlobalHandles env
-> Var Extent2D
forall b c a. (b -> c) -> (a -> b) -> a -> c
. App GlobalHandles env -> 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 :: Lens' (App GlobalHandles st, Frame rp p rr) (SomeRef st)
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
  { forall rp p rr st. Stage rp p rr st -> Text
sTitle :: Text

  , forall rp p rr st.
Stage rp p rr st
-> SwapchainResources -> ResourceT (StageRIO st) rp
sAllocateRP :: SwapchainResources -> ResourceT (StageRIO st) rp
  , forall rp p rr st.
Stage rp p rr st
-> SwapchainResources -> rp -> ResourceT (StageRIO st) p
sAllocateP  :: SwapchainResources -> rp -> ResourceT (StageRIO st) p
  , forall rp p rr st.
Stage rp p rr st
-> StageRIO (Maybe SwapchainResources) (ReleaseKey, st)
sInitialRS  :: StageRIO (Maybe SwapchainResources) (ReleaseKey, st)
  , forall rp p rr 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
  , forall rp p rr st.
Stage rp p rr st -> st -> rr -> StageFrameRIO rp p rr st ()
sUpdateBuffers    :: st -> rr -> StageFrameRIO rp p rr st ()
  , forall 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
  { forall renderpass pipelines resources.
Frame renderpass pipelines resources -> Word64
fIndex                       :: Word64 -- ^ Which number frame is this
  , forall renderpass pipelines resources.
Frame renderpass pipelines resources -> Window
fWindow                      :: Window
  , forall renderpass pipelines resources.
Frame renderpass pipelines resources -> SurfaceKHR
fSurface                     :: Khr.SurfaceKHR

  , forall renderpass pipelines resources.
Frame renderpass pipelines resources -> SwapchainResources
fSwapchainResources          :: SwapchainResources
  , forall renderpass pipelines resources.
Frame renderpass pipelines resources -> renderpass
fRenderpass                  :: renderpass
  , forall renderpass pipelines resources.
Frame renderpass pipelines resources -> pipelines
fPipelines                   :: pipelines
  , forall renderpass pipelines resources.
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.
    -}

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

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

  , forall renderpass pipelines resources.
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.
    -}

  , forall renderpass pipelines resources.
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
  { forall a. RecycledResources a -> Semaphore
rrImageAvailableSemaphore :: Vk.Semaphore
    -- ^ A binary semaphore passed to 'acquireNextImageKHR'
  , forall a. RecycledResources a -> Semaphore
rrRenderFinishedSemaphore :: Vk.Semaphore
    -- ^ A binary semaphore to synchronize rendering and presenting

  , forall a. 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)
    -}

  , forall a. RecycledResources a -> a
rrData                    :: a
  }

instance HasLogFunc env => HasLogFunc (env, Frame rp p rr) where
  logFuncL :: Lens' (env, Frame rp p rr) LogFunc
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 :: forall a. 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