module Render.ImGui
  ( allocate
  -- , allocateWithFonts
  , allocateWith

  , beforeLoop
  , afterLoop

  , capturingKeyboard
  , capturingMouse

  , mkDrawData
  , draw
  ) where

import RIO

import Control.Monad.Trans.Resource (ReleaseKey, ResourceT, register, release)
import DearImGui.Vulkan (InitInfo(..))
import DearImGui.Vulkan qualified as ImGui
import DearImGui qualified as ImGui
import DearImGui.GLFW (glfwNewFrame, glfwShutdown)
import DearImGui.GLFW.Vulkan (glfwInitForVulkan)
import RIO.App (appEnv)
import Vulkan.Core10 qualified as Vk
import Vulkan.Exception (VulkanException(..))
import Vulkan.Utils.QueueAssignment (QueueFamilyIndex(..))

import Engine.Types (GlobalHandles(..), StageRIO)
import Engine.Vulkan.Types (HasRenderPass(..), HasSwapchain(..), Queues(..), getMultisample)
import Resource.CommandBuffer qualified as CommandBuffer
import Resource.DescriptorSet qualified as DescriptorSet

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 (FilePath, Float)
--   -> ResourceT (StageRIO st) (ReleaseKey, t (Maybe ImGui.Font))
-- allocateWithFonts swapchain renderpass subpassIx fonts =
--   allocateWith swapchain renderpass subpassIx do
--     loaded <- for fonts \(path, size) ->
--       ImGui.addFontFromFileTTF path size
--     ImGui.buildFontAtlas
--     _atlasKey <- register ImGui.clearFontAtlas
--     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
  _ctxKey <- register $ ImGui.destroyContext ctx

  context@GlobalHandles{..} <- asks appEnv

  let (QueueFamilyIndex queueFamily, queue) = qGraphics ghQueues

  (_poolKey, pool) <- DescriptorSet.allocatePool 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 :: DescriptorSet.TypeMap 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
  ]

{- | 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 "Setting up ImGui"
  window <- asks $ ghWindow . appEnv
  success <- glfwInitForVulkan window installCallbacks
  unless success do
    logWarn "glfwInitForVulkan failed"

afterLoop :: StageRIO st ()
afterLoop = do
  logDebug "Shutting down ImGui"
  glfwShutdown

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)