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)
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
  { ghWindow             :: GLFW.Window
  , ghSurface            :: Khr.SurfaceKHR
  , ghInstance           :: Vk.Instance
  , ghPhysicalDevice     :: Vk.PhysicalDevice
  , ghPhysicalDeviceInfo :: Vulkan.PhysicalDeviceInfo
  , ghDevice             :: Vk.Device
  , ghAllocator          :: VMA.Allocator
  , ghQueues             :: Vulkan.Queues (QueueFamilyIndex, Vk.Queue)
  , ghScreenP            :: ProjectionProcess
  , ghStageSwitch        :: StageSwitchVar
  }

getScreenP :: MonadReader (App GlobalHandles s) m => m ProjectionProcess
getScreenP = asks $ ghScreenP . appEnv

instance HasVulkan GlobalHandles where
  getInstance           = ghInstance
  getQueues             = ghQueues
  getPhysicalDevice     = ghPhysicalDevice
  getPhysicalDeviceInfo = ghPhysicalDeviceInfo
  getDevice             = ghDevice
  getAllocator          = 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)

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

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

  , sBeforeLoop       :: StageRIO st a
  , sUpdateBuffers    :: st -> rr -> 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
  { fIndex                       :: Word64 -- ^ Which number frame is this
  , fWindow                      :: Window
  , fSurface                     :: Khr.SurfaceKHR

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

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

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

  , 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.
    -}

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

  , 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)
    -}

  , rrData                    :: a
  }

instance HasLogFunc env => HasLogFunc (env, Frame rp p rr) where
  logFuncL = _1 . logFuncL

instance MonadResource (RIO (env, Frame rp p rr)) where
  {-# INLINE liftResourceT #-}
  liftResourceT rt =
    asks (snd . fResources . snd) >>=
      liftIO . ResourceT.runInternalState rt