{-# LANGUAGE AllowAmbiguousTypes #-}

module Resource.Texture
  ( Texture(..)
  , destroy

  , TextureError(..)
    -- * Texture types
  , Flat
  , CubeMap
  , ArrayOf
  , TextureLayers(..)

    -- * Utilities
  , allocateCollectionWith
  , allocateTextureWith
  , debugNameCollection
  , TextureLoader

  , createImageView
  , imageCI
  , imageAllocationCI
  , stageBufferCI
  , stageAllocationCI
  ) where

import RIO

import Data.Bits ((.|.))
import Data.List qualified as List
import GHC.Stack (withFrozenCallStack)
import GHC.TypeLits (Nat, KnownNat, natVal)
import RIO.FilePath (takeBaseName)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.NamedType ((:::))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)
import VulkanMemoryAllocator qualified as VMA

import Engine.Vulkan.Types (HasVulkan(getDevice), MonadVulkan, Queues)
import Resource.Collection qualified as Collection
import Resource.Image (AllocatedImage(..), destroyAllocatedImage, subresource)
import Resource.Image qualified as Image

data TextureError
  = LoadError Int64 Text
  | LayerError Word32 Word32
  | MipLevelsError Word32 Int
  | ArrayError Word32 Word32
  deriving (Eq, Ord, Show)

instance Exception TextureError

data Texture a = Texture
  { tFormat         :: Vk.Format
  , tMipLevels      :: Word32
  , tLayers         :: Word32 -- ^ Actual number of layers, up to @ArrayOf a@
  , tAllocatedImage :: AllocatedImage
  }
  deriving (Show)

data CubeMap
data Flat
data ArrayOf (layers :: Nat)

-- | Number of expected texture layers to load from resource.
class TextureLayers a where
  textureLayers :: Word32

instance TextureLayers CubeMap where
  textureLayers = 6

instance TextureLayers Flat where
  textureLayers = 1

instance KnownNat n => TextureLayers (ArrayOf n) where
  textureLayers = fromInteger $ natVal (Proxy @n)

type TextureLoader m layers = Vk.Format -> Queues Vk.CommandPool -> FilePath -> m (Texture layers)

type TextureLoaderAction m layers = FilePath -> m (Texture layers)

-- * Allocation wrappers

allocateCollectionWith
  :: (Resource.MonadResource m, MonadVulkan env m, Traversable t)
  => TextureLoaderAction m layers
  -> t FilePath
  -> m (Resource.ReleaseKey, t (Texture layers))
allocateCollectionWith action collection = do
  res <- traverse (allocateTextureWith action) collection
  key <- Resource.register $
    traverse_ (Resource.release . fst) res
  pure (key, fmap snd res)

allocateTextureWith
  :: (Resource.MonadResource m, MonadVulkan env m)
  => TextureLoaderAction m layers
  -> FilePath
  -> m (Resource.ReleaseKey, Texture layers)
allocateTextureWith action path = do
  context <- ask
  createTexture <- toIO $ action path
  Resource.allocate createTexture (destroy context)

debugNameCollection
  :: ( Traversable t
     , MonadVulkan env m
     , HasLogFunc env
     , HasCallStack
     )
  => t (Texture layers)
  -> t FilePath
  -> m ()
debugNameCollection textures paths = do
  device <- asks getDevice
  for_ names \((ix, path), Texture{tAllocatedImage}) -> do
    withFrozenCallStack . logDebug $ displayShow (ix, path)
    Debug.nameObject device (Image.aiImage tAllocatedImage) $
      fromString $ show @Natural ix <> ":" <> takeBaseName path
  where
    names = List.zip
      (toList $ Collection.enumerate paths)
      (toList textures)

-- * Implementation

destroy :: (MonadIO io, HasVulkan context) => context -> Texture a -> io ()
destroy context Texture{tAllocatedImage} =
  destroyAllocatedImage context tAllocatedImage

createImageView
  :: (MonadIO io, HasVulkan context)
  => context
  -> Vk.Image
  -> Vk.Format
  -> "mip levels" ::: Word32
  -> "array layers" ::: Word32
  -> io Vk.ImageView
createImageView context image format mipLevels arrayLayers =
  Vk.createImageView (getDevice context) imageViewCI Nothing
  where
    imageViewCI = zero
      { Vk.image            = image
      , Vk.viewType         = viewType
      , Vk.format           = format
      , Vk.components       = zero
      , Vk.subresourceRange = colorRange
      }

    viewType =
      if arrayLayers == 6 then
        Vk.IMAGE_VIEW_TYPE_CUBE
      else
        Vk.IMAGE_VIEW_TYPE_2D

    colorRange =
      subresource Vk.IMAGE_ASPECT_COLOR_BIT mipLevels arrayLayers

imageCI :: Vk.Format -> Vk.Extent3D -> Word32 -> Word32 -> Vk.ImageCreateInfo '[]
imageCI format extent mipLevels arrayLayers = zero
  { Vk.flags         = flags
  , Vk.imageType     = Vk.IMAGE_TYPE_2D
  , Vk.format        = format
  , Vk.extent        = extent
  , Vk.mipLevels     = mipLevels
  , Vk.arrayLayers   = if isCube then 6 else arrayLayers
  , Vk.tiling        = Vk.IMAGE_TILING_OPTIMAL
  , Vk.initialLayout = Vk.IMAGE_LAYOUT_UNDEFINED
  , Vk.usage         = usage
  , Vk.sharingMode   = Vk.SHARING_MODE_EXCLUSIVE
  , Vk.samples       = Vk.SAMPLE_COUNT_1_BIT -- XXX: no multisampling here
  }
  where
    isCube =
      arrayLayers == 6

    usage =
      Vk.IMAGE_USAGE_SAMPLED_BIT .|.  -- Sampler
      Vk.IMAGE_USAGE_TRANSFER_DST_BIT -- Staging

    flags =
      if isCube then
        Vk.IMAGE_CREATE_CUBE_COMPATIBLE_BIT
      else
        zero

imageAllocationCI :: VMA.AllocationCreateInfo
imageAllocationCI = zero
  { VMA.usage         = VMA.MEMORY_USAGE_GPU_ONLY
  , VMA.requiredFlags = Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT
  }

stageBufferCI :: Integral a => a -> Vk.BufferCreateInfo '[]
stageBufferCI pixelBytes = zero
  { Vk.size        = fromIntegral pixelBytes
  , Vk.usage       = Vk.BUFFER_USAGE_TRANSFER_SRC_BIT
  , Vk.sharingMode = Vk.SHARING_MODE_EXCLUSIVE
  }

stageAllocationCI :: VMA.AllocationCreateInfo
stageAllocationCI = zero
  { VMA.flags         = VMA.ALLOCATION_CREATE_MAPPED_BIT
  , VMA.usage         = VMA.MEMORY_USAGE_CPU_TO_GPU
  , VMA.requiredFlags = Vk.MEMORY_PROPERTY_HOST_VISIBLE_BIT
  }