module Render.ImGui
  ( allocate

  , allocateWithFonts
  , allocateWith

  , renderWith

  , allocateLoop
  , beforeLoop
  , afterLoop

  , capturingKeyboard
  , capturingMouse

  , mkDrawData
  , draw
  ) where

import RIO

import Control.Monad.Trans.Resource (ReleaseKey, ResourceT, register, release)
import DearImGui qualified as ImGui
import DearImGui.FontAtlas qualified as FontAtlas
import DearImGui.GLFW (glfwNewFrame, glfwShutdown)
import DearImGui.GLFW.Vulkan (glfwInitForVulkan)
import DearImGui.Vulkan (InitInfo(..))
import DearImGui.Vulkan qualified as ImGui
import Engine.Stage.Component qualified as Stage
import Engine.Types (GlobalHandles(..), StageRIO)
import Engine.Vulkan.Types (HasRenderPass(..), HasSwapchain(..), Queues(..), getDevice, getMultisample)
import Resource.CommandBuffer qualified as CommandBuffer
import Resource.Region qualified as Region
import Resource.Vulkan.DescriptorPool qualified as DescriptorPool
import RIO.App (appEnv)
import Vulkan.Core10 qualified as Vk
import Vulkan.Exception (VulkanException(..))
import Vulkan.NamedType (type (:::))
import Vulkan.Utils.QueueAssignment (QueueFamilyIndex(..))

allocate
  :: ( HasSwapchain swapchain
     , HasRenderPass renderpass
     )
  => swapchain
  -> renderpass
  -> Word32
  -> ResourceT (StageRIO st) ReleaseKey
allocate swapchain renderpass subpassIx =
  fmap fst $
    allocateWith swapchain renderpass subpassIx $
      pure ()

allocateWithFonts
  :: ( HasSwapchain swapchain
     , HasRenderPass renderpass
     , Traversable t
     )
  => swapchain
  -> renderpass
  -> Word32
  -> t FontAtlas.FontSource
  -> ResourceT (StageRIO st) (ReleaseKey, t ImGui.Font)
allocateWithFonts swapchain renderpass subpassIx fonts =
  allocateWith swapchain renderpass subpassIx do
    loaded <- FontAtlas.rebuild fonts
    _atlasKey <- register FontAtlas.clear
    pure loaded

allocateWith
  :: ( HasSwapchain swapchain
     , HasRenderPass renderpass
     )
  => swapchain
  -> renderpass
  -> Word32
  -> ResourceT (StageRIO st) a
  -> ResourceT (StageRIO st) (ReleaseKey, a)
allocateWith swapchain renderpass subpassIx action = do
  logDebug "Initializing DearImGui"
  debugReleaseFinished <- toIO (logDebug "Released DearImGui")
  void $! register debugReleaseFinished

  ctx <- ImGui.createContext
  Region.register_ $ ImGui.destroyContext ctx

  context@GlobalHandles{..} <- asks appEnv

  let (QueueFamilyIndex queueFamily, queue) = qGraphics ghQueues

  pool <- Region.local $
    DescriptorPool.allocate (Just "ImGui") 1 dsSizes

  let
    initInfo = InitInfo
      { instance'      = ghInstance
      , physicalDevice = ghPhysicalDevice
      , device         = ghDevice
      , queueFamily    = queueFamily
      , queue          = queue

      , minImageCount  = getMinImageCount swapchain
      , imageCount     = getImageCount swapchain
      , msaaSamples    = getMultisample swapchain

      , subpass        = subpassIx
      , pipelineCache  = Vk.NULL_HANDLE

      , descriptorPool = pool

      , mbAllocator    = Nothing
      , checkResult    = \case { Vk.SUCCESS -> pure (); e -> throwM $ VulkanException e }
      }

  res@(_cb, initOk) <- ImGui.vulkanInit initInfo (getRenderPass renderpass)
  unless initOk do
    logError "DearImGui vulkan initialization failed"
    exitFailure
  key <- register $ ImGui.vulkanShutdown res

  actionRes <- action

  -- TODO: oneshots
  (oneshotPoolKey, oneshotPool) <- CommandBuffer.allocatePools context
  CommandBuffer.oneshot_ context oneshotPool qGraphics \cb -> do
    fontsOk <- ImGui.vulkanCreateFontsTexture cb
    unless fontsOk do
      logError "ImGui.vulkanCreateFontsTexture failed"
      exitFailure
  release oneshotPoolKey
  ImGui.vulkanDestroyFontUploadObjects

  debugReleaseStart <- toIO (logDebug "Releasing DearImGui")
  void $! register debugReleaseStart

  pure (key, actionRes)

dsSizes :: [(Vk.DescriptorType, Word32)]
dsSizes = map (, 100)
  [ Vk.DESCRIPTOR_TYPE_SAMPLER
  , Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
  , Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
  , Vk.DESCRIPTOR_TYPE_STORAGE_IMAGE
  , Vk.DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER
  , Vk.DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER
  , Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER
  , Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC
  , Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC
  , Vk.DESCRIPTOR_TYPE_INPUT_ATTACHMENT
  ]

allocateLoop :: Bool -> ResourceT (StageRIO st) ()
allocateLoop installCallbacks = do
  lift $ beforeLoop installCallbacks
  shutdownImGui <- lift $ toIO afterLoop
  void $! register shutdownImGui

{- | Initialize context to serve the draws in the current render loop.

You will need to trigger callbacks from DearImGui.GLFW if you're opting out of them here.

DearImgui will run your previously installed GLFW callbacks.
-}
beforeLoop :: Bool -> StageRIO st ()
beforeLoop installCallbacks = do
  logDebug "glfwInitForVulkan"
  window <- asks $ ghWindow . appEnv
  success <- glfwInitForVulkan window installCallbacks
  unless success do
    logWarn "glfwInitForVulkan failed"

afterLoop :: StageRIO st ()
afterLoop = do
  logDebug "glfwShutdown"
  asks getDevice >>= Vk.deviceWaitIdle
  logDebug "glfwShutdown: Device idle"
  glfwShutdown
  logDebug "glfwShutdown finished"

mkDrawData :: MonadIO m => m a -> m (a, ImGui.DrawData)
mkDrawData action = do
  ImGui.vulkanNewFrame
  glfwNewFrame
  ImGui.newFrame
  result <- action
  ImGui.render
  drawData <- ImGui.getDrawData
  pure (result, drawData)

draw :: MonadIO m => ImGui.DrawData -> Vk.CommandBuffer -> m ()
draw drawData commandBuffer = do
  ImGui.vulkanRenderDrawData drawData commandBuffer Nothing

capturingKeyboard :: MonadIO m => m () -> m ()
capturingKeyboard action =
  ImGui.wantCaptureKeyboard >>= (`unless` action)

capturingMouse :: MonadIO m => m () -> m ()
capturingMouse action =
  ImGui.wantCaptureMouse >>= (`unless` action)

renderWith
  :: HasRenderPass renderpass
  => (t -> renderpass)
  -> "subpass index" ::: Word32
  -> Stage.Rendering t p st
  -> Stage.Rendering t p st
renderWith getRP subpassIx Stage.Rendering{..} = Stage.Rendering
  { rAllocateRP
  , rAllocateP = \swapchain rps -> do
      void $! allocate swapchain (getRP rps) subpassIx
      rAllocateP swapchain rps
  }