{-# 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 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.Model qualified as Model
import Resource.Region qualified as Region
import Resource.Source (Source)
import Resource.Texture.Ktx2 qualified as Ktx2
import RIO.State (gets)
import RIO.Vector.Partial ((!))
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
  -> (Font.Config, Font.Config)
  -> (Source, Source)
  -> ((Text -> StageSetupRIO ()) -> StageSetupRIO loaded)
  -> (loaded -> StackStage)
  -> ( Setup Vector Vector loaded -> Engine.StackStage
     , Engine.StageSetupRIO (Setup Vector Vector loaded)
     )
bootstrap titleMessage (smallFont, largeFont) (bgPath, spinnerPath) loadAction nextStage =
  (stackStageBootstrap, action)
  where
    action = withPools \pools -> do
      logDebug "Bootstrapping loader"

      fonts <- traverse (Font.allocate pools) [smallFont, largeFont]
      textures <- traverse (Ktx2.load pools) [bgPath, spinnerPath]

      let
        fontContainers = fmap Font.container fonts
        combinedTextures = Collection.enumerate CombinedTextures.Collection
          { textures = textures
          , fonts    = fmap Font.texture fonts
          }

      let
        uiSettings = UI.Settings
          { titleMessage = titleMessage
          , backgroundIx = 0
          , spinnerIx    = 1

          , combined = combinedTextures
          , fonts    = fontContainers

          , smallFont = \fs -> fs ! 0
          , largeFont = \fs -> fs ! 1
          }

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

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

stackStageBootstrap
  :: (Traversable fonts, Traversable textures)
  => Setup fonts textures loaded -> StackStage
stackStageBootstrap Setup{..} = stackStage loadAction nextStage uiSettings

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

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

    allocateRenderPass swapchain = do
      Basic.allocate
        Basic.Settings
          { sShadowLayers = 1
          , sShadowSize   = 1
          }
        swapchain

    allocatePipelines swapchain rps = do
      logDebug "Allocating loader pipelines"
      samplers <- Samplers.allocate (Swapchain.getAnisotropy swapchain)
      Basic.allocatePipelines
        (sceneBinds samplers)
        (Swapchain.getMultisample swapchain)
        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 textures fonts
  -> 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

    pure RunState{..}

initialFrameResources
  :: (Traversable fonts, Traversable textures)
  => UI.Settings fonts textures
  -> 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{..}