module Engine.Frame
  ( Frame(..)
  , initial
  , run
  , advance
  , queueSubmit

  , RecycledResources(..)
  , initialRecycledResources
  , timeoutError
  ) where

import RIO

import Control.Monad.Trans.Resource (ResourceT, MonadResource, allocate, release)
import Control.Monad.Trans.Resource qualified as ResourceT
import GHC.IO.Exception (IOErrorType(TimeExpired), IOException(IOError))
import RIO.App (appEnv)
import RIO.Text qualified as Text
import RIO.Vector qualified as Vector
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Core12.Promoted_From_VK_KHR_timeline_semaphore qualified as Vk12
import Vulkan.CStruct.Extends (SomeStruct(..), pattern (:&), pattern (::&))
import Vulkan.NamedType ((:::))
import Vulkan.Utils.QueueAssignment (QueueFamilyIndex(..))
import Vulkan.Zero (zero)

import Engine.DataRecycler (DumpResource, WaitResource)
import Engine.Setup.Window qualified as Window
import Engine.Types (GlobalHandles(..), StageRIO, Stage(..), Frame(..), GPUWork, RecycledResources(..))
import Engine.Types.Options (optionsPresent, optionsMsaa)
import Engine.Types.RefCounted (newRefCounted)
import Engine.Vulkan.Swapchain (SwapchainResources(..), SwapchainInfo(..), allocSwapchainResources, recreateSwapchainResources)
import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, RenderPass(..), Queues)

initial
  :: Maybe SwapchainResources
  -> DumpResource (RecycledResources rr)
  -> Stage rp p rr st
  -> StageRIO st (Frame rp p rr)
initial oldSR dumpResource Stage{..} = do
  logDebug "Making initial frame"

  GlobalHandles{..} <- asks appEnv
  let device = ghDevice

  let
    fPresent = optionsPresent ghOptions
    fMSAA = optionsMsaa ghOptions

  sfSwapchainResources <- case oldSR of
    Nothing -> do
      windowSize <- liftIO $ Window.getExtent2D ghWindow
      let oldSwapchain = Vk.NULL_HANDLE
      allocSwapchainResources
        fPresent
        fMSAA
        oldSwapchain
        windowSize
        ghSurface
        ghScreenVar
    Just old ->
      pure old

  {- XXX:
    Create this resource object at the global level so it's closed correctly on exception.
  -}
  (stageKey, stageResources) <- allocate ResourceT.createInternalState ResourceT.closeInternalState
  stageRefCounted <- newRefCounted $ release stageKey
  semiFrame <- flip ResourceT.runInternalState stageResources do
    {- XXX:
      Stages appearing on the top of the stage stack are to create their swapchain-derived resources.
      Don't keep the release keys, all resources here are refcounted and live for the lifetime of the stage.
      Resources will be released when the stage is finished or suspended and all the frames are done.
    -}

    debugAlloc <- toIO . logDebug $ "Allocating inside stage " <> display sTitle
    debugRelease <- toIO . logDebug $ "Releasing inside stage " <> display sTitle
    void $! ResourceT.allocate_ debugAlloc debugRelease

    -- For each render pass:
    sfRenderpass <- sAllocateRP sfSwapchainResources

    -- TODO: Recreate this if the swapchain format changes
    sfPipelines <- sAllocateP sfSwapchainResources sfRenderpass

    (_, sfRenderFinishedHostSemaphore) <- Vk.withSemaphore
      device
      (zero ::& Vk12.SemaphoreTypeCreateInfo Vk12.SEMAPHORE_TYPE_TIMELINE 0 :& ())
      Nothing
      allocate

    logDebug $ "Creating initial recycled resources for stage " <> display sTitle
    sfRecycledResources <- initialRecycledResources sInitialRR sfRenderpass sfPipelines
    replicateM_ (INFLIGHT_FRAMES - 1) do
      resources <- initialRecycledResources sInitialRR sfRenderpass sfPipelines
      liftIO $ dumpResource resources

    releaseDataDebug <- toIO . logDebug $ "Releasing recycled resources for stage " <> display sTitle
    void $! Resource.register releaseDataDebug

    pure
      ( sfSwapchainResources
      , sfRenderpass
      , sfPipelines
      , sfRenderFinishedHostSemaphore
      , sfRecycledResources
      )

  let
    (fSwapchainResources, fRenderpass, fPipelines, fRenderFinishedHostSemaphore, fRecycledResources) = semiFrame

  {- XXX:
    Create this resource object at the global level so it's closed correctly on exception.
    Recycled frame resources can linger for a bit longer after its stage is gone, thus 'RefCounted'.
  -}
  fResources <- allocate ResourceT.createInternalState ResourceT.closeInternalState

  fGPUWork <- liftIO $ newIORef mempty

  pure Frame
    { fIndex          = 1
    , fWindow         = ghWindow
    , fSurface        = ghSurface
    , fStageResources = (stageRefCounted, stageResources)
    , ..
    }

pattern INFLIGHT_FRAMES :: (Eq a, Num a) => a
pattern INFLIGHT_FRAMES = 2 -- XXX: up to two frames submitted for rendering

-- | Derive next frame
advance
  :: ( HasLogFunc env
     , HasVulkan env
     , MonadResource (RIO env)
     , RenderPass rp
     )
  => WaitResource (RecycledResources rr)
  -> Frame rp p rr
  -> Bool
  -> RIO env (Frame rp p rr)
advance waitDumped f needsNewSwapchain = do
  -- Wait for a prior frame to finish, then we can steal it's resources!

  -- Handle mvar indefinite timeout exception here:
  -- https://github.com/expipiplus1/vulkan/issues/236
  fRecycledResources <-
    liftIO waitDumped >>= \case
      Left block ->
        timeout 15e6 (liftIO block) >>= \case
          Nothing -> do
            logError . display $ Text.unwords
              [ "Timed out waiting for recycled resources."
              , "A recycler thread is stuck on timeline semaphore or something."
              , "Try running with --recycler-wait 15000 or a similar value."
              ]
            exitFailure
          Just rr ->
            pure rr
      Right rs ->
        pure rs

  (fSwapchainResources, fRenderpass) <- getNext f

  -- The per-frame resource helpers need to be created fresh
  fGPUWork   <- liftIO $ newIORef mempty
  fResources <- allocate ResourceT.createInternalState ResourceT.closeInternalState

  pure Frame
    { fIndex                       = fIndex f + 1
    , fWindow                      = fWindow f
    , fSurface                     = fSurface f
    , fPipelines                   = fPipelines f
    , fRenderFinishedHostSemaphore = fRenderFinishedHostSemaphore f
    , fStageResources              = fStageResources f
    , fPresent                     = fPresent f
    , fMSAA                        = fMSAA f
    , fSwapchainResources
    , fRenderpass
    , fGPUWork
    , fResources
    , fRecycledResources
    }
  where
    getNext Frame{..} = do
      if needsNewSwapchain then do
        windowSize <- liftIO $ Window.getExtent2D fWindow
        newResources <- recreateSwapchainResources fPresent fMSAA windowSize fSwapchainResources

        let
          formatMatch =
            siSurfaceFormat (srInfo newResources) ==
            siSurfaceFormat (srInfo fSwapchainResources)
        unless formatMatch do
          logWarn "Swapchain changed format"
          throwString "TODO: Handle swapchain changing formats"

        newRenderpass <- updateRenderpass newResources fRenderpass

        pure
          ( newResources
          , newRenderpass
          )
      else
        pure
          ( fSwapchainResources
          , fRenderpass
          )

run
  :: ( HasLogFunc env
     , HasVulkan env
     , MonadResource (RIO env)
     )
  => (RecycledResources rr -> IO ())
  -> Maybe Int
  -> RIO (env, Frame rp p rr) a
  -> Frame rp p rr
  -> RIO env a
run recycle recyclerWait render frame@Frame{..} = do
  env <- ask
  runRIO (env, frame) render `finally` void (spawn flush)
  where
    flush = do
      device <- asks getDevice
      waits <- readIORef fGPUWork
      let tenSecondsKhr = 10e9
      -- logDebug $ "Waiting Frame " <> displayShow fIndex

      unless (null waits) do
        traverse_ threadDelay recyclerWait
        let
          waitInfo = zero
            { Vk12.semaphores = Vector.fromList (map fst waits)
            , Vk12.values     = Vector.fromList (map snd waits)
            }
        waitTwice waitInfo tenSecondsKhr >>= \case
          Vk.TIMEOUT -> do
            logError "Time out (10s) waiting for frame to finish on Device"
            timeoutError "Time out (10s) waiting for frame to finish on Device"
            {-
              XXX: recycler thread will crash now, never recycling its resources,
              resulting in an indefinite MVar block.
            -}
          Vk.SUCCESS ->
            pure ()
          huh ->
            logWarn $ "waitTwice returned " <> displayShow huh

      -- logDebug $ "Flushing Frame " <> displayShow fIndex

      -- Free resources wanted elsewhere now, all those in RecycledResources
      for_ (rrQueues fRecycledResources) \commandPool ->
        Vk.resetCommandPool device commandPool Vk.COMMAND_POOL_RESET_RELEASE_RESOURCES_BIT

      -- Signal we're done by making the recycled resources available
      liftIO $ recycle fRecycledResources

      -- Destroy frame-specific resources at our leisure
      release (fst fResources)

-- | 'queueSubmit' and add wait for the timeline 'Semaphore' before retiring the frame.
queueSubmit
  :: MonadVulkan env m
  => Vk.Queue
  -> Vector (SomeStruct Vk.SubmitInfo)
  -> IORef [GPUWork]
  -> Vk.Semaphore
  -> Word64
  -> m ()
queueSubmit q submits gpuWork hostSemaphore frameIndex = do
  {-
    Make sure we don't get interrupted between submitting the work and
    recording the wait.
  -}
  mask \_ -> do
    Vk.queueSubmit q submits Vk.NULL_HANDLE
    atomicModifyIORef' gpuWork \waits ->
      ( (hostSemaphore, frameIndex) : waits
      , ()
      )

initialRecycledResources
  :: ( Resource.MonadResource (RIO env)
     , HasVulkan env
     , HasLogFunc env
     )
  => (Queues Vk.CommandPool -> rp -> p -> ResourceT (RIO env) rr)
  -> rp
  -> p
  -> ResourceT (RIO env) (RecycledResources rr)
initialRecycledResources initialRecycledData rps pipes = do
  device <- asks getDevice

  (_iaKey, rrImageAvailableSemaphore) <- Vk.withSemaphore
    device
    (zero ::& Vk12.SemaphoreTypeCreateInfo Vk12.SEMAPHORE_TYPE_BINARY 0 :& ())
    Nothing
    allocate

  (_rfKey, rrRenderFinishedSemaphore) <- Vk.withSemaphore
    device
    (zero ::& Vk12.SemaphoreTypeCreateInfo Vk12.SEMAPHORE_TYPE_BINARY 0 :& ())
    Nothing
    allocate

  queues <- asks getQueues
  rrQueues <- for queues \(QueueFamilyIndex ix, _queue) -> do
    let
      commandPoolCI = Vk.CommandPoolCreateInfo
        { flags            = zero
        , queueFamilyIndex = ix
        }
    cpDebug <- toIO . logDebug $ "Release time for command pool for queue " <> display ix
    void $! ResourceT.register cpDebug
    fmap snd $! Vk.withCommandPool device commandPoolCI Nothing ResourceT.allocate

  rrData <- initialRecycledData rrQueues rps pipes

  pure RecycledResources{..}

{- |
  Wait for some semaphores, if the wait times out give the frame one last
  chance to complete with a zero timeout.

  It could be that the program was suspended during the preceding
  wait causing it to timeout, this will check if it actually
  finished.
-}
waitTwice
  :: (MonadVulkan env m, HasLogFunc env)
  => Vk12.SemaphoreWaitInfo
  -> "timeout" ::: Word64
  -> m Vk.Result
waitTwice waitInfo t = do
  device <- asks getDevice
  Vk12.waitSemaphoresSafe device waitInfo t >>= \case
    Vk.TIMEOUT -> do
      r <- Vk12.waitSemaphoresSafe device waitInfo 1e3
      logWarn $ mconcat
        [ "waiting a second time on " <> displayShow waitInfo
        , " got " <> displayShow r
        ]
      pure r
    r ->
      pure r

timeoutError :: MonadThrow m => String -> m a
timeoutError message =
  throwM $ IOError Nothing TimeExpired "" message Nothing Nothing

spawn :: (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
spawn action = do
  actionIO <- toIO action
  {-
    If we don't remove the release key when the thread is done it'll leak,
    remove it at the end of the async action when the thread is going to
    die anyway.

    Mask this so there's no chance we're inturrupted before writing the mvar.
  -}
  kv  <- newEmptyMVar
  mask $ \_ -> do
    (k, r) <- allocate
      (asyncWithUnmask \unmask ->
        unmask $ actionIO <* (Resource.unprotect =<< liftIO (readMVar kv))
      )
      uninterruptibleCancel
    putMVar kv k
    pure r