{-# LANGUAGE AllowAmbiguousTypes #-}

module Resource.Texture
  ( Texture(..)

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

    -- * Utilities
  , debugNameCollection
  , TextureLoader

  , createImageView
  , imageCI
  , imageAllocationCI
  , stageBufferCI
  , stageAllocationCI

  , withSize2d
  , withSize3d
  ) where

import RIO

import Data.Bits ((.|.))
import Data.List qualified as List
import Geomancy (Vec2, vec2)
import GHC.Stack (withFrozenCallStack)
import GHC.TypeLits (Nat, KnownNat, natVal)
import GHC.Records (HasField(..))
import RIO.FilePath (takeBaseName)
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(..))
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 tag = 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)

-- * Allocation wrappers

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

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 =
      Image.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
  }

-- * Helpers

{-# INLINE withSize2d #-}
withSize2d :: Num i => (i -> i -> a) -> Texture tag -> a
withSize2d f t =
  f
    (fromIntegral width)
    (fromIntegral height)
  where
    Vk.Extent3D{width, height} =
      Image.aiExtent (tAllocatedImage t)

{-# INLINE withSize3d #-}
withSize3d :: Num i => (i -> i -> i -> a) -> Texture tag -> a
withSize3d f t =
  f
    (fromIntegral width)
    (fromIntegral height)
    (fromIntegral depth)
  where
    Vk.Extent3D{width, height, depth} =
      Image.aiExtent (tAllocatedImage t)

instance HasField "size" (Texture tag) Vec2 where
  {-# INLINE getField #-}
  getField = withSize2d vec2