module Engine.Types where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Control.Monad.Trans.Resource qualified as ResourceT
import Data.Kind (Type)
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.Options (Options)
import Engine.Types.RefCounted (RefCounted)
import Engine.Vulkan.Swapchain (SwapchainResources(..))
import Engine.Vulkan.Types (HasVulkan(..), HasSwapchain(..))
import Engine.Vulkan.Types qualified as Vulkan
import Engine.Worker qualified as Worker

-- * App globals

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

{-# INLINE askScreenVar #-}
askScreenVar
  :: MonadReader (App GlobalHandles st) m
  => m (Worker.Var Vk.Extent2D)
askScreenVar = asks $ ghScreenVar . 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)

instance HasStateRef st (App GlobalHandles st, Frame rp p rr) where
  stateRefL =
    lens
      (appState . fst)
      (\(app, frame) st' ->
          ( app
              { appState = st'
              }
          , frame
          )
      )

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

  , sAllocateRP :: SwapchainResources -> ResourceT (StageRIO st) rp
  , 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

  , fPresent :: Maybe Khr.PresentModeKHR
  , fMSAA    :: Vk.SampleCountFlagBits

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

instance HasSwapchain (Frame renderpass pipelines resources) where
  getSurfaceExtent  = getSurfaceExtent . fSwapchainResources
  getSurfaceFormat  = getSurfaceFormat . fSwapchainResources
  getDepthFormat    = getDepthFormat . fSwapchainResources
  getMultisample    = getMultisample . fSwapchainResources
  getAnisotropy     = getAnisotropy . fSwapchainResources
  getSwapchainViews = getSwapchainViews . fSwapchainResources
  getMinImageCount  = getMinImageCount . fSwapchainResources
  getImageCount     = getImageCount . fSwapchainResources
  {-# INLINE getSurfaceExtent #-}
  {-# INLINE getSurfaceFormat #-}
  {-# INLINE getDepthFormat #-}
  {-# INLINE getMultisample #-}
  {-# INLINE getAnisotropy #-}
  {-# INLINE getSwapchainViews #-}
  {-# INLINE getMinImageCount #-}
  {-# INLINE getImageCount #-}

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

type HKD :: (Type -> Type) -> Type -> Type
type family HKD f a where
  HKD Identity a = a
  HKD f a = f a