{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedRecordDot #-}

module Engine.Vulkan.Swapchain
  ( SwapchainResources(..)
  , SwapchainInfo(..)
  , allocSwapchainResources
  , recreateSwapchainResources
  , createSwapchain
  , threwSwapchainError

  , HasSwapchain(..)

  , setDynamic
  , setDynamicFullscreen
  ) where

import RIO

import Data.Bits (zeroBits, (.&.), (.|.))
import RIO.Vector qualified as V
import UnliftIO.Resource (MonadResource, allocate, release)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Exception (VulkanException(..))
import Vulkan.Extensions.VK_KHR_surface qualified as Khr
import Vulkan.Extensions.VK_KHR_swapchain qualified as Khr
import Vulkan.Utils.Misc ((.&&.))
import Vulkan.NamedType ((:::))
import Vulkan.Zero (zero)

import Engine.Types.RefCounted (RefCounted, newRefCounted, releaseRefCounted)
import Engine.Vulkan.Types (MonadVulkan, HasVulkan(..), HasSwapchain(..), pdiProperties)
import Engine.Worker qualified as Worker

data SwapchainResources = SwapchainResources
  { srInfo       :: SwapchainInfo
  , srImageViews :: Vector Vk.ImageView
  , srImages     :: Vector Vk.Image
  , srRelease    :: RefCounted
  , srScreenVar  :: Worker.Var Vk.Extent2D
  }

data SwapchainInfo = SwapchainInfo
  { siSwapchain           :: Khr.SwapchainKHR
  , siSwapchainReleaseKey :: Resource.ReleaseKey
  , siPresentMode         :: Khr.PresentModeKHR
  , siMinImageCount       :: Word32
  , siSurfaceFormat       :: Vk.Format
  , siSurfaceColorspace   :: Khr.ColorSpaceKHR
  , siDepthFormat         :: Vk.Format
  , siMultisample         :: Vk.SampleCountFlagBits
  , siAnisotropy          :: Float
  , siImageExtent         :: Vk.Extent2D
  , siSurface             :: Khr.SurfaceKHR
  }

instance HasSwapchain SwapchainResources where
  getSurfaceExtent  = siImageExtent . srInfo
  getSurfaceFormat  = siSurfaceFormat . srInfo
  getDepthFormat    = siDepthFormat . srInfo
  getMultisample    = siMultisample . srInfo
  getAnisotropy     = siAnisotropy . srInfo
  getSwapchainViews = srImageViews
  getMinImageCount  = siMinImageCount . srInfo
  getImageCount     = fromIntegral . V.length . srImages

-- | Allocate everything which depends on the swapchain
allocSwapchainResources
  :: ( MonadResource (RIO env)
     , HasVulkan env
     , HasLogFunc env
     )
  => Maybe Khr.PresentModeKHR
  -> Vk.SampleCountFlagBits
  -> Khr.SwapchainKHR
  -- ^ Previous swapchain, can be NULL_HANDLE
  -> Vk.Extent2D
  -- ^ If the swapchain size determines the surface size, use this size
  -> Khr.SurfaceKHR
  -> Worker.Var Vk.Extent2D
  -> RIO env SwapchainResources
  -- -> ResourceT (RIO env) SwapchainResources
allocSwapchainResources present msaa oldSwapchain windowSize surface screenVar = do
  logDebug "Allocating swapchain resources"

  device <- asks getDevice
  info@SwapchainInfo{..} <- createSwapchain present msaa oldSwapchain windowSize surface

  -- XXX: Get all the swapchain images, and create views for them
  (_, swapchainImages) <- Khr.getSwapchainImagesKHR device siSwapchain
  res <- for swapchainImages $
    createImageView siSurfaceFormat
  let (imageViewKeys, imageViews) = V.unzip res

  -- XXX: This refcount is released in 'recreateSwapchainResources'
  releaseDebug <- toIO $ logDebug "Releasing swapchain resources"
  releaseResources <- newRefCounted $ do
    releaseDebug
    traverse_ release imageViewKeys
    release siSwapchainReleaseKey

  Worker.pushInput screenVar $ const windowSize

  pure SwapchainResources
    { srInfo       = info
    , srImageViews = imageViews
    , srImages     = swapchainImages
    , srRelease    = releaseResources
    , srScreenVar  = screenVar
    }

recreateSwapchainResources
  :: ( MonadResource (RIO env)
     , HasVulkan env
     , HasLogFunc env
     )
  => Maybe Khr.PresentModeKHR
  -> Vk.SampleCountFlagBits
  -> Vk.Extent2D
  -> SwapchainResources
  -- ^ The reference to these resources will be dropped
  -> RIO env SwapchainResources
recreateSwapchainResources present msaa windowSize oldResources = do
  sr <- allocSwapchainResources
    present
    msaa
    (siSwapchain $ srInfo oldResources)
    windowSize
    (siSurface $ srInfo oldResources)
    (srScreenVar oldResources)
  releaseRefCounted (srRelease oldResources)
  pure sr

-- | Create a swapchain from a 'SurfaceKHR'
createSwapchain
  :: ( MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     )
  => Maybe Khr.PresentModeKHR
  -> Vk.SampleCountFlagBits
  -> Khr.SwapchainKHR
  -- ^ Old swapchain, can be NULL_HANDLE
  -> Vk.Extent2D
  -- ^ If the swapchain size determines the surface size, use this size
  -> Khr.SurfaceKHR
  -> m SwapchainInfo
createSwapchain present msaa oldSwapchain explicitSize surf = do
  physical <- asks getPhysicalDevice
  device <- asks getDevice
  props <- asks $ pdiProperties . getPhysicalDeviceInfo
  surfaceCaps <- Khr.getPhysicalDeviceSurfaceCapabilitiesKHR physical surf

  logDebug $ displayShow surfaceCaps

  -- Check flags
  for_ requiredUsageFlags \flag ->
    unless (Khr.supportedUsageFlags surfaceCaps .&&. flag) do
      logError $ "Surface images do not support " <> displayShow flag
      throwString $ "Surface images do not support " <> show flag

  -- Select a present mode
  desiredPresentModes <- case present of
    Nothing -> do
      logDebug $ "Using default present modes: " <> displayShow defaultPresentModes
      pure defaultPresentModes
    Just selected -> do
      logDebug $ "Forcing selected present mode: " <> displayShow selected
      pure [selected]
  (_, availablePresentModes) <- Khr.getPhysicalDeviceSurfacePresentModesKHR physical surf
  presentMode <-
    case filter (`V.elem` availablePresentModes) desiredPresentModes of
      [] -> do
        logError "Unable to find a suitable present mode for swapchain"
        throwString "Unable to find a suitable present mode for swapchain"
      x : _rest ->
        pure x

  -- Select a surface format
  -- getPhysicalDeviceSurfaceFormatsKHR doesn't return an empty list
  -- (_, availableFormats) <- Khr.getPhysicalDeviceSurfaceFormatsKHR physical surf
  surfaceFormatKhr <- getSurfaceFormatKhr physical surf preferSrgb

  depthFormat <- getDepthFormats physical preferStenciledDepth >>= \case
    fmt : _rest ->
      pure fmt
    _none ->
      throwString "Unable to find a suitable depth format"

  -- Calculate the extent
  let
    imageExtent =
      case Khr.currentExtent (surfaceCaps :: Khr.SurfaceCapabilitiesKHR) of
        Vk.Extent2D w h | w == maxBound, h == maxBound ->
          explicitSize
        extent ->
          extent

  let
    minImageCount =
      let
        limit = case surfaceCaps.maxImageCount of
          0 -> maxBound
          n -> n
        -- Request one additional image to prevent us having to wait for
        -- the driver to finish
        buffer = 1
        desired =
          buffer + surfaceCaps.minImageCount
      in
        min limit desired

  let compositeAlphaMode = Khr.COMPOSITE_ALPHA_OPAQUE_BIT_KHR
  unless (compositeAlphaMode .&&. Khr.supportedCompositeAlpha surfaceCaps) $
    throwString $ "Surface doesn't support " <> show compositeAlphaMode

  let
    Khr.SurfaceFormatKHR{colorSpace=surfaceColorspace, format=surfaceFormat} = surfaceFormatKhr
    swapchainCreateInfo = Khr.SwapchainCreateInfoKHR
      { surface            = surf
      , next               = ()
      , flags              = zero
      , queueFamilyIndices = mempty -- No need to specify when not using concurrent access
      , minImageCount      = minImageCount
      , imageFormat        = surfaceFormat
      , imageColorSpace    = surfaceColorspace
      , imageExtent        = imageExtent
      , imageArrayLayers   = 1
      , imageUsage         = foldr (.|.) zero requiredUsageFlags
      , imageSharingMode   = Vk.SHARING_MODE_EXCLUSIVE
      , preTransform       = Khr.currentTransform (surfaceCaps :: Khr.SurfaceCapabilitiesKHR)
      , compositeAlpha     = compositeAlphaMode
      , presentMode        = presentMode
      , clipped            = True
      , oldSwapchain       = oldSwapchain
      }

  logDebug $ "Creating swapchain from " <> displayShow swapchainCreateInfo
  (key, swapchain) <- Khr.withSwapchainKHR device swapchainCreateInfo Nothing allocate

  pure SwapchainInfo
    { siSwapchain           = swapchain
    , siSwapchainReleaseKey = key
    , siPresentMode         = presentMode
    , siMinImageCount       = minImageCount
    , siSurface             = surf
    , siSurfaceFormat       = surfaceFormat
    , siSurfaceColorspace   = surfaceColorspace
    , siDepthFormat         = depthFormat
    , siMultisample         = msaaSamples msaa props
    , siAnisotropy          = Vk.maxSamplerAnisotropy (Vk.limits props)
    , siImageExtent         = imageExtent
    }

-- -- The vector passed will have at least one element
-- selectSurfaceFormat :: Vector Khr.SurfaceFormatKHR -> Khr.SurfaceFormatKHR
-- selectSurfaceFormat = V.maximumBy (comparing surfaceFormatScore)
--   where
--     -- An ordered list of formats to choose for the swapchain images, if none
--     -- match then the first available format will be chosen.
--     surfaceFormatScore :: Khr.SurfaceFormatKHR -> Int
--     surfaceFormatScore = \case
--       _ -> 0

getSurfaceFormatKhr
  :: MonadIO io
  => Vk.PhysicalDevice
  -> Khr.SurfaceKHR
  -> Khr.SurfaceFormatKHR
  -> io Khr.SurfaceFormatKHR
getSurfaceFormatKhr device surface desiredFormat = do
  (_res, formats) <- Khr.getPhysicalDeviceSurfaceFormatsKHR device surface
  pure case toList formats of
    [] ->
      desiredFormat
    [Khr.SurfaceFormatKHR Vk.FORMAT_UNDEFINED _colorSpace] ->
      desiredFormat
    candidates | any cond candidates ->
      desiredFormat
    whatever : _rest ->
      whatever
  where
    cond f =
      Khr.format f == Khr.format desiredFormat &&
      Khr.colorSpace f == Khr.colorSpace desiredFormat

preferSrgb :: Khr.SurfaceFormatKHR
preferSrgb =
  Khr.SurfaceFormatKHR
    Vk.FORMAT_B8G8R8A8_SRGB
    Khr.COLOR_SPACE_SRGB_NONLINEAR_KHR

getDepthFormats
  :: MonadIO io
  => Vk.PhysicalDevice
  -> [Vk.Format]
  -> io [Vk.Format]
getDepthFormats device desiredDepthFormats = do
  properties <- traverse (Vk.getPhysicalDeviceFormatProperties device) desiredDepthFormats
  pure do
    (format, props) <- zip desiredDepthFormats properties
    guard $
      Vk.optimalTilingFeatures props .&&. Vk.FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT
    pure format

preferStenciledDepth :: [Vk.Format]
preferStenciledDepth =
  [ Vk.FORMAT_D32_SFLOAT_S8_UINT
  , Vk.FORMAT_D24_UNORM_S8_UINT
  , Vk.FORMAT_D32_SFLOAT
  ]

msaaSamples
  :: Vk.SampleCountFlagBits
  -> Vk.PhysicalDeviceProperties
  -> Vk.SampleCountFlagBits
msaaSamples upperLimit Vk.PhysicalDeviceProperties{limits} =
  case samplesAvailable of
    [] ->
      -- XXX: Something went wrong...
      Vk.SAMPLE_COUNT_1_BIT
    best : _rest ->
      best
  where
    counts =
      Vk.framebufferColorSampleCounts limits .&.
      Vk.framebufferDepthSampleCounts limits

    samplesAvailable = do
      countBit <- msaaCandidates
      guard $ countBit <= upperLimit -- XXX: requested
      guard $ (counts .&. countBit) /= zeroBits -- XXX: capable
      pure countBit

msaaCandidates :: [Vk.SampleCountFlagBits]
msaaCandidates =
  [ Vk.SAMPLE_COUNT_64_BIT -- XXX: extremely unrealistic?
  , Vk.SAMPLE_COUNT_32_BIT -- XXX: unrealistic?
  , Vk.SAMPLE_COUNT_16_BIT -- XXX: possible, but not that impactful
  , Vk.SAMPLE_COUNT_8_BIT
  , Vk.SAMPLE_COUNT_4_BIT -- XXX: Khronos-recommended for rasterizing
  , Vk.SAMPLE_COUNT_2_BIT
  , Vk.SAMPLE_COUNT_1_BIT -- XXX: i.e. "disable", always available
  ]

-- | An ordered list of the present mode to be chosen for the swapchain.
defaultPresentModes :: [Khr.PresentModeKHR]
defaultPresentModes =
  [ Khr.PRESENT_MODE_FIFO_RELAXED_KHR
  , Khr.PRESENT_MODE_FIFO_KHR --  ^ This will always be present
  ]

-- | The images in the swapchain must support these flags.
requiredUsageFlags :: [Vk.ImageUsageFlagBits]
requiredUsageFlags =
  [ Vk.IMAGE_USAGE_COLOR_ATTACHMENT_BIT
  -- , Vk.IMAGE_USAGE_STORAGE_BIT
  ]

-- | Catch an ERROR_OUT_OF_DATE_KHR exception and return 'True' if that happened
threwSwapchainError :: MonadUnliftIO f => f () -> f Bool
threwSwapchainError = fmap isLeft . tryJust swapchainError
 where
  swapchainError = \case
    VulkanException e@Vk.ERROR_OUT_OF_DATE_KHR ->
      Just e

    VulkanException Vk.ERROR_SURFACE_LOST_KHR ->
      error "TODO: handle ERROR_SURFACE_LOST_KHR"

    VulkanException _ ->
      Nothing

-- | Create a pretty vanilla ImageView covering the whole image
createImageView
  :: ( MonadResource m
     , MonadVulkan env m
     )
  => Vk.Format
  -> Vk.Image
  -> m (Resource.ReleaseKey, Vk.ImageView)
createImageView format image = do
  device <- asks getDevice
  Vk.withImageView device imageViewCI Nothing allocate
 where
  imageViewCI = zero
    { Vk.image = image
    , Vk.viewType         = Vk.IMAGE_VIEW_TYPE_2D
    , Vk.format           = format
    -- , Vk.components       = zero
    , Vk.subresourceRange = zero
        { Vk.aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
        , Vk.baseMipLevel   = 0
        , Vk.levelCount     = 1
        , Vk.baseArrayLayer = 0
        , Vk.layerCount     = 1
        }
    }

setDynamic
  :: MonadIO io
  => Vk.CommandBuffer
  -> "viewport" ::: Vk.Rect2D
  -> "scissor"  ::: Vk.Rect2D
  -> io ()
setDynamic cb viewrect scissor = do
  Vk.cmdSetViewport cb 0 [viewport]
  Vk.cmdSetScissor cb 0 [scissor]
  where
    viewport = Vk.Viewport
      { x        = realToFrac x
      , y        = realToFrac y
      , width    = realToFrac width
      , height   = realToFrac height
      , minDepth = 0
      , maxDepth = 1
      }
      where
        Vk.Offset2D{x, y} = offset
        Vk.Extent2D{width, height} = extent
        Vk.Rect2D{offset, extent} = viewrect

setDynamicFullscreen
  :: ( HasSwapchain swapchain
     , MonadIO io
     )
  => Vk.CommandBuffer
  -> swapchain
  -> io ()
setDynamicFullscreen cb swapchain = do
  Vk.cmdSetViewport cb
    0
    [ Vk.Viewport
        { x        = 0
        , y        = 0
        , width    = realToFrac width
        , height   = realToFrac height
        , minDepth = 0
        , maxDepth = 1
        }
    ]
  Vk.cmdSetScissor cb 0
    [ Vk.Rect2D
        { offset = zero
        , extent = extent
        }
    ]
  where
    extent@Vk.Extent2D{width, height} =
      getSurfaceExtent swapchain