{-# LANGUAGE OverloadedLists #-}

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

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Engine.StageSwitch (trySwitchStage)
import Engine.Types (StackStage(..), Stage(..), StageSetupRIO)
import Engine.Types qualified as Engine
import Engine.UI.Layout qualified as Layout
import Engine.UI.Message qualified as Message
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.Source (Source)
import Resource.Texture qualified as Texture
import Resource.Texture.Ktx1 qualified as Ktx1
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.Types (FrameResources(..), RunState(..))
import Stage.Loader.Scene qualified as Scene
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"

      let
        fontConfigs = [smallFont, largeFont] :: Vector Font.Config
      (fontKey, fonts) <- Font.allocateCollection pools fontConfigs

      let
        texturePaths = [bgPath, spinnerPath] :: Vector Source
      (textureKey, textures) <- Texture.allocateCollectionWith
        (Ktx1.load pools)
        texturePaths

      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
          }

      loaderKey <- Resource.register $
        traverse_ @[] Resource.release
          [ fontKey
          , textureKey
          ]

      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
  , loaderKey  :: Resource.ReleaseKey
  }

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
  { sTitle = "Loader"

  , sAllocateP  = allocatePipelines
  , sInitialRS  = initialRunState loadAction nextStage uiSettings
  , sInitialRR  = initialFrameResources uiSettings
  , sBeforeLoop = pure ()

  , sUpdateBuffers  = Render.updateBuffers
  , sRecordCommands = Render.recordCommands

  , sAfterLoop    = pure
  }
  where
    allocatePipelines swapchain rps = do
      logDebug "Allocating loader pipelines"
      (_, samplers) <- Samplers.allocate swapchain
      Basic.allocatePipelines (sceneBinds samplers) swapchain rps

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

initialRunState
  :: ((Text -> StageSetupRIO ()) -> StageSetupRIO loaded)
  -> (loaded -> StackStage)
  -> UI.Settings textures fonts
  -> StageSetupRIO (Resource.ReleaseKey, RunState)
initialRunState loadAction nextStage uiSettings =
  withPools \pools -> do
    rsProjectionP <- Engine.getScreenP

    (sceneUiKey, rsSceneUiP) <- Worker.registered $
      Scene.spawn rsProjectionP

    context <- ask

    rsQuadUV <- Model.createStagedL context pools (Quad.toVertices Quad.texturedQuad) Nothing
    quadKey <- Resource.register $ Model.destroyIndexed context rsQuadUV

    (screenKey, screenBoxP) <- Worker.registered Layout.trackScreen

    (uiKey, rsUI) <- UI.spawn pools screenBoxP uiSettings

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

    releaseKeys <- Resource.register $
      traverse_ @[] Resource.release
        [ sceneUiKey
        , quadKey
        , screenKey
        , uiKey
        ]

    switcher <- 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 (releaseKeys, 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{..}