module Engine.Run
  ( runStack
  , run
  , step
  ) where

import RIO

import Graphics.UI.GLFW qualified as GLFW
import RIO.App (App(..), appEnv)
import RIO.State (get, put)
import UnliftIO.Resource (ReleaseKey, release)
import Vulkan.Core10 qualified as Vk
import GHC.Clock (getMonotonicTimeNSec)

import Engine.DataRecycler (DataRecycler(..))
import Engine.DataRecycler qualified as DataRecycler
import Engine.Frame qualified as Frame
import Engine.Render (renderFrame)
import Engine.StageSwitch (getNextStage)
import Engine.Types (Frame, NextStage(..), RecycledResources, StageStack, StackStage(..), Stage(..), StageRIO)
import Engine.Types qualified as Engine
import Engine.Types.Options (optionsRecyclerWait, optionsMaxFPS)
import Engine.Types.RefCounted (releaseRefCounted)
import Engine.Vulkan.Swapchain (SwapchainResources(srRelease), threwSwapchainError)
import Engine.Vulkan.Types (RenderPass, getDevice)

runStack :: StageStack -> StageRIO (Maybe SwapchainResources) ()
runStack = \case
  [] -> do
    logInfo "Stage stack finished"
    get >>= \case
      Nothing ->
        pure ()
      Just oldSR ->
        releaseRefCounted (srRelease oldSR)
    asks getDevice >>= Vk.deviceWaitIdle
    threadDelay 0.5e6 -- XXX: this is arbitrary and may not always satisfy the validator

  StackStage stage : rest -> do
    logInfo $ "Setting up stage " <> display (sTitle stage)
    (stageRelease, runner) <- prepareStage (sInitialRS stage)
    oldSR <- get
    runner (run oldSR stageRelease stage) >>= proceed stageRelease rest

  StackStageContinue stKey state stage : rest -> do
    logInfo $ "Resuming stage " <> display (sTitle stage)
    (stageRelease, runner) <- prepareStage (pure (stKey, state))
    oldSR <- get
    runner (run oldSR stageRelease stage) >>= proceed stageRelease rest

type StageResult = (SwapchainResources, StageAction)

proceed :: ReleaseKey -> StageStack -> StageResult -> StageRIO (Maybe SwapchainResources) ()
proceed stageRelease rest (oldSR, stageAction) = do
  put (Just oldSR)
  case stageAction of
    StageDone -> do
      releaseStage "Stage done"
      runStack rest
    StageReplace nextStage -> do
      releaseStage "Stage replaced"
      runStack (nextStage : rest)
    StagePush frozenStage nextStage -> do
      case frozenStage of
        StackStage{} ->
          logDebug "Restarting stage pushed"
        StackStageContinue{} ->
          logDebug "Frozen stage pushed"
      runStack (nextStage : frozenStage : rest)
  where
    releaseStage label = do
      logDebug label
      void $! async do
        {-
          XXX: wait for the current render job to finish.

          Since the new frames wouldn't be submitted until the control flow
          gets back to the render loop, this should be enough of a signal that
          the stage resources aren't used anymore.

          The call is costly, but its effects aren't seen.
        -}
        asks getDevice >>= Vk.deviceWaitIdle
        release stageRelease
        logDebug $ label <> ", released"

prepareStage
  :: StageRIO env (ReleaseKey, st)
  -> StageRIO env (ReleaseKey, StageRIO st a -> StageRIO env a)
prepareStage initialRS = do
  (key, rs) <- initialRS
  freshStateVar <- newSomeRef rs
  App{..} <- ask
  let
    stageApp = App
      { appState = freshStateVar
      , ..
      }
  pure (key, runRIO stageApp)

run
  :: RenderPass rp
  => Maybe SwapchainResources
  -> ReleaseKey
  -> Stage rp p rr st
  -> StageRIO st StageResult
run oldSR stKey stage@Stage{..} = do
  logInfo $ "Starting stage: " <> display sTitle

  recycler <- DataRecycler.new

  startFrame <- Frame.initial oldSR (drDump recycler) stage

  Engine.GlobalHandles{ghWindow, ghOptions} <- asks appEnv
  quit <- liftIO $ GLFW.windowShouldClose ghWindow
  if quit then do
    logDebug $ "Forcing stage unwind for " <> display sTitle
    pure
      ( Frame.fSwapchainResources startFrame
      , StageDone
      )
  else do
    logInfo $ "Entering stage loop: " <> display sTitle
    startTime <- getMonotonicTime
    (finalFrame, stageAction) <- bracket sBeforeLoop sAfterLoop \_stagePrivates ->
      stageLoop (step stKey stage recycler (optionsMaxFPS ghOptions)) startFrame
    endTime <- getMonotonicTime
    let
      frames  = Frame.fIndex finalFrame
      seconds = endTime - startTime

    logInfo $ "Stage finished: " <> display sTitle
    logInfo $ "Running time: " <> display seconds
    logInfo $ "Average FPS: " <> display (fromIntegral frames / seconds)
    releaseRefCounted $ fst (Frame.fStageResources finalFrame)
    pure
      ( Frame.fSwapchainResources finalFrame
      , stageAction
      )

step
  :: RenderPass rp
  => ReleaseKey
  -> Stage rp p rr st
  -> DataRecycler (RecycledResources rr)
  -> Maybe Int
  -> Frame rp p rr
  -> StageRIO st (LoopAction (Frame rp p rr))
step stKey stage@Stage{..} DataRecycler{..} maxFPSM frame = do
  startTime <- liftIO getMonotonicTimeNSec
  liftIO GLFW.pollEvents

  -- XXX: hard unwind all the stages, rendering nothing
  quit <- liftIO $ GLFW.windowShouldClose (Frame.fWindow frame)

  if quit then
    pure LoopQuit
  else
    getNextStage >>= \case
      Nothing -> do
        needsNewSwapchain <- threwSwapchainError do
          Engine.GlobalHandles{ghOptions} <- asks appEnv
          let recyclerWait = optionsRecyclerWait ghOptions
          rs <- get
          Frame.run drDump recyclerWait (renderFrame (sUpdateBuffers rs) sRecordCommands) frame
        nextFrame <- Frame.advance drWait frame needsNewSwapchain

        for_ maxFPSM \maxFPS -> do
          -- Wait until the end of allocated time.
          endTime <- liftIO getMonotonicTimeNSec
          let elapsedUS = (fromEnum $ endTime - startTime) `div` 1e3 :: Int
              fpsCapUS = 1e6 `div` maxFPS
              waitUS = fpsCapUS - min fpsCapUS elapsedUS
          threadDelay waitUS

        pure $ LoopNextFrame nextFrame
      Just Finish ->
        -- XXX: finish the stage and proceed with the remaining stack
        pure LoopQuit
      Just (Replace nextStage) ->
        pure $ LoopReplaceStage frame nextStage
      Just (PushRestart nextStage) ->
        pure $ LoopPushStage frame (StackStage stage) nextStage
      Just (PushFreeze nextStage) -> do
        frozen <- freeze stKey stage
        pure $ LoopPushStage frame frozen nextStage

freeze :: RenderPass rp => ReleaseKey -> Stage rp p rr st -> StageRIO st StackStage
freeze stKey stage = do
  st <- get
  pure $ StackStageContinue stKey st stage

data StageAction
  = StageReplace StackStage
  | StagePush StackStage StackStage
  | StageDone

data LoopAction f
  = LoopNextFrame f
  | LoopReplaceStage f StackStage
  | LoopPushStage f StackStage StackStage
  | LoopQuit

{-# INLINE stageLoop #-}
stageLoop :: (f -> StageRIO st (LoopAction f)) -> f -> StageRIO st (f, StageAction)
stageLoop action current =
  action current >>= \case
    LoopNextFrame nextFrame ->
      stageLoop action nextFrame
    LoopQuit ->
      pure (current, StageDone)
    LoopReplaceStage lastFrame nextStage ->
      pure (lastFrame, StageReplace nextStage)
    LoopPushStage lastFrame frozenStage nextStage ->
      pure (lastFrame, StagePush frozenStage nextStage)