Setup.hs

{-# LANGUAGE OverloadedLists #-}

module Stage.Loader.Setup
  ( bootstrap
  , stackStage
  ) where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Engine.Camera qualified as Camera
import Engine.Stage.Component qualified as Stage
import Engine.StageSwitch (trySwitchStage)
import Engine.Types (StackStage(..), StageSetupRIO)
import Engine.Types qualified as Engine
import Engine.UI.Message qualified as Message
import Engine.Vulkan.Swapchain qualified as Swapchain
import Engine.Vulkan.Types (Queues)
import Engine.Worker qualified as Worker
import Geometry.Quad qualified as Quad
import RIO.State (gets)
import Render.Basic qualified as Basic
import Render.DescSets.Set0 qualified as Set0
import Render.Samplers qualified as Samplers
import Resource.Collection qualified as Collection
import Resource.Combined.Textures qualified as CombinedTextures
import Resource.CommandBuffer (withPools)
import Resource.Font qualified as Font
import Resource.Font.Ktxf qualified as Ktxf
import Resource.Model qualified as Model
import Resource.Region qualified as Region
import Resource.Source (Source)
import Resource.Texture.Ktx2 qualified as Ktx2
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk

import Stage.Loader.Render qualified as Render
import Stage.Loader.Scene qualified as Scene
import Stage.Loader.Types (FrameResources(..), RunState(..))
import Stage.Loader.UI qualified as UI

bootstrap
  :: Text
  -> (Source, Source)
  -> (Source, Source)
  -> ((Text -> StageSetupRIO ()) -> StageSetupRIO loaded)
  -> (loaded -> StackStage)
  -> ( Setup loaded -> Engine.StackStage
     , Engine.StageSetupRIO (Setup loaded)
     )
bootstrap titleMessage (smallFontSrc, largeFontSrc) (bgPath, spinnerPath) loadAction nextStage =
  (stackStageBootstrap, action)
  where
    action = withPools \pools -> do
      logDebug "Bootstrapping loader"

      smallFont <- Font.allocate pools smallFontSrc
      largeFont <- Font.allocate pools largeFontSrc

      textures <- traverse (Ktx2.load pools) [bgPath, spinnerPath]

      let
        combinedTextures = Collection.enumerate CombinedTextures.Collection
          { textures = textures
          , fonts    = [smallFont.texture, largeFont.texture]
          }

      let
        (backgroundIx, spinnerIx, smallFontId, largeFontId) =
          case map fst $ toList combinedTextures of
            [a, b, c, d] -> (a, b, c, d)
            _ -> error "assert: combinedTextures has exactly 4 elements"

      smallFontStack <- Ktxf.allocateStack1 smallFontId smallFont.bundle
      largeFontStack <- Ktxf.allocateStack1 largeFontId largeFont.bundle

      let
        uiSettings = UI.Settings
          { titleMessage = titleMessage
          , backgroundIx = backgroundIx
          , spinnerIx    = spinnerIx
          , combined = combinedTextures
          , smallFont = smallFontStack
          , largeFont = largeFontStack
          , fontBundles = [smallFont.bundle, largeFont.bundle]
          }

      logDebug "Finished bootstrapping loader"
      pure Setup{..}

data Setup loaded = Setup
  { loadAction :: (Text -> StageSetupRIO ()) -> StageSetupRIO loaded
  , nextStage  :: loaded -> StackStage
  , uiSettings :: UI.Settings
  }

stackStageBootstrap :: Setup loaded -> StackStage
stackStageBootstrap Setup{..} = stackStage loadAction nextStage uiSettings

stackStage
  :: ((Text -> StageSetupRIO ()) -> StageSetupRIO loaded)
  -> (loaded -> StackStage)
  -> UI.Settings
  -> StackStage
stackStage loadAction nextStage uiSettings =
  StackStage $ loaderStage loadAction nextStage uiSettings

loaderStage
  :: ((Text -> StageSetupRIO ()) -> StageSetupRIO loaded)
  -> (loaded -> StackStage)
  -> UI.Settings
  -> Basic.Stage FrameResources RunState
loaderStage loadAction nextStage uiSettings = Stage.assemble "Loader" rendering resources (Just scene)
  where
    rendering = Stage.Rendering
      { rAllocateRP = Basic.allocate_
      , rAllocateP = allocatePipelines
      }

    allocatePipelines swapchain rps = do
      logDebug "Allocating loader pipelines"
      samplers <- Samplers.allocate (Swapchain.getAnisotropy swapchain)
      Basic.allocatePipelines
        (Basic.swapchainSettings swapchain)
        (sceneBinds samplers)
        rps

    sceneBinds samplers = Set0.mkBindings
      samplers
      (UI.combined uiSettings)
      Nothing
      0

    resources = Stage.Resources
      { rInitialRS = initialRunState loadAction nextStage uiSettings
      , rInitialRR = initialFrameResources uiSettings
      }

    scene = Stage.Scene
      { scBeforeLoop = pure ()
      , scUpdateBuffers = Render.updateBuffers
      , scRecordCommands = Render.recordCommands
      }

initialRunState
  :: ((Text -> StageSetupRIO ()) -> StageSetupRIO loaded)
  -> (loaded -> StackStage)
  -> UI.Settings
  -> StageSetupRIO (Resource.ReleaseKey, RunState)
initialRunState loadAction nextStage uiSettings =
  withPools \pools -> Region.run do
    rsProjectionP <- Camera.spawnOrthoPixelsCentered

    rsSceneUiP <- Scene.spawn rsProjectionP

    rsQuadUV <- Model.createStagedL (Just "rsQuadUV") pools (Quad.toVertices Quad.texturedQuad) Nothing
    Model.registerIndexed_ rsQuadUV

    rsUI <- Region.local $
      UI.spawn pools uiSettings

    let
      updateProgress text = do
        logInfo $ "Loader: " <> display text
        Worker.pushInput (UI.progressInput rsUI) \msg -> msg
          { Message.inputText = text
          }

    switcher <- lift $ async do
      loader <- async do
        logDebug "Starting load action"
        try (loadAction updateProgress) >>= \case
          Left (e :: SomeException) -> do
            logError $ "Load action failed with " <> displayShow e
            throwM e
          Right r -> do
            logDebug "Load action finished"
            pure r

      link loader
      -- threadDelay 1e6
      waitCatch loader >>= \case
        Left oopsie -> do
          logError "Loader failed"
          throwM oopsie
        Right loaded -> do
          logInfo "Loader signalled a stage change"

          updateProgress "Done!"
          switched <- trySwitchStage . Engine.Replace $
            nextStage loaded
          unless switched $
            logError "Loader switch failed"

    -- XXX: propagate exceptions from loader threads
    link switcher

    -- transfer the references
    let rsFontAssets = uiSettings.fontBundles

    pure RunState{..}

initialFrameResources
  :: UI.Settings
  -> Queues Vk.CommandPool
  -> Basic.RenderPasses
  -> Basic.Pipelines
  -> ResourceT (Engine.StageRIO RunState) FrameResources
initialFrameResources UI.Settings{combined} _pools _passes pipelines = do
  frSceneUi <- Set0.allocate
    (Basic.getSceneLayout pipelines)
    (fmap snd combined)
    Nothing
    Nothing
    mempty -- XXX: no shadows on loader
    Nothing

  frUI <- gets rsUI >>=
    UI.newObserver

  pure FrameResources{..}