module Resource.Texture.Ktx1
  ( load
  , loadBytes
  , loadKtx1
  ) where

import RIO

import Codec.Ktx qualified as Ktx1
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Data.Text qualified as Text
import Data.Vector qualified as Vector
import Foreign qualified
import GHC.Stack (withFrozenCallStack)
import Vulkan.Core10 qualified as Vk
import Vulkan.Utils.FromGL qualified as FromGL
import VulkanMemoryAllocator qualified as VMA

import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, Queues)
import Resource.Image (AllocatedImage(..))
import Resource.Image qualified as Image
import Resource.Source (Source(..))
import Resource.Source qualified as Source
import Resource.Texture (Texture(..), TextureLayers(..))
import Resource.Texture qualified as Texture

load
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadThrow m
     , HasLogFunc env
     , Typeable a
     , HasCallStack
     )
  => Queues Vk.CommandPool
  -> Source
  -> m (Texture a)
load pool source =
  withFrozenCallStack $
    Source.load (loadBytes pool) source

loadBytes
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> ByteString
  -> m (Texture a)
loadBytes pool bytes = do
  case Ktx1.fromByteString bytes of
    Left (offset, err) -> do
      logError $ "Texture load error: " <> fromString err
      throwM $ Texture.LoadError offset (Text.pack err)
    Right ktx1 ->
      loadKtx1 pool ktx1

loadKtx1
  :: forall a m env
  .  ( TextureLayers a
     , MonadVulkan env m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> Ktx1.Ktx
  -> m (Texture a)
loadKtx1 pool Ktx1.Ktx{header, images=ktxImages} = do
  context <- ask
  let vma = getAllocator context

  let
    skipMips = 0 -- DEBUG: Vector.length ktxImages `div` 2
    mipsSkipped = min (Vector.length ktxImages - 1) skipMips
    images = Vector.drop mipsSkipped ktxImages
    mipLevels = Ktx1.numberOfMipmapLevels header - fromIntegral mipsSkipped

  when (Vector.null images) $
    throwM $ Texture.LoadError 0 "At least one image must be present in KTX"

  unless (fromIntegral mipLevels == Vector.length images) $
    throwM $ Texture.MipLevelsError mipLevels (Vector.length images)

  -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vk_format.h#L676
  format <- case Ktx1.glInternalFormat header of
    -- XXX: Force all BC7s to SRGB
    36492 ->
      pure Vk.FORMAT_BC7_SRGB_BLOCK
    other ->
      case FromGL.internalFormat other of
        Nothing ->
          error $ "Unexpected glInternalFormat: " <> show other -- TODO: throwIo
        Just fmt ->
          -- XXX: going in blind
          pure fmt
  logDebug $ mconcat
    [ "Loading format "
    , display $ Ktx1.glInternalFormat header
    , " as "
    , displayShow format
    ]

  -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vkloader.c#L552
  let
    extent = Vk.Extent3D
      { Vk.width  = Ktx1.pixelWidth header `Foreign.shiftR` mipsSkipped
      , Vk.height = Ktx1.pixelHeight header `Foreign.shiftR` mipsSkipped
      , Vk.depth  = max 1 $ Ktx1.pixelDepth header
      }
    arrayLayers = max 1 $ Ktx1.numberOfArrayElements header

  -- TODO: basisu can encode movies as arrays, this could be handy
  unless (arrayLayers == 1) do
    logError "TODO: arrayLayers > 1"
    throwM $ Texture.ArrayError 1 arrayLayers

  let
    numLayers = Ktx1.numberOfFaces header
    mipSizes = fmap ((*) numLayers . Ktx1.imageSize) images

    offsets' = Vector.scanl' (+) 0 mipSizes
    totalSize = Vector.last offsets'
    offsets = Vector.init offsets'

  logDebug $ "mipSizes: " <> displayShow mipSizes
  logDebug $ "offsets: " <> displayShow offsets

  {- XXX:
    Image created before staging buffer due to `Image.copyBufferToImage`
    issued inside `VMA.withBuffer` bracket.
  -}
  (image, allocation, _info) <- VMA.createImage
    vma
    (Texture.imageCI format extent mipLevels numLayers)
    Texture.imageAllocationCI

  Image.transitionLayout
    context
    pool
    image
    mipLevels
    numLayers -- XXX: arrayLayers is always 0 for now
    format
    Vk.IMAGE_LAYOUT_UNDEFINED
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL

  unless (numLayers == textureLayers @a) $
    throwM $ Texture.ArrayError (textureLayers @a) numLayers

  VMA.withBuffer vma (Texture.stageBufferCI totalSize) Texture.stageAllocationCI bracket \(staging, stage, stageInfo) -> do
    let ixMipImages = Vector.zip3 (Vector.fromList [0..]) offsets images
    Vector.forM_ ixMipImages \(mipIx, offset, Ktx1.MipLevel{imageSize, arrayElements}) -> do
      let ixArrayElements = Vector.zip (Vector.fromList [0..]) arrayElements
      Vector.forM_ ixArrayElements \(arrayIx, Ktx1.ArrayElement{faces}) -> do
        let ixFaces = Vector.zip (Vector.fromList [0..]) faces
        Vector.forM_ ixFaces \(faceIx, Ktx1.Face{zSlices}) -> do
          let ixSlices = Vector.zip (Vector.fromList [0..]) zSlices
          Vector.forM_ ixSlices \(sliceIx, Ktx1.ZSlice{block}) -> do
            let
              indices = mconcat
                [ "["
                , " mip:" <> display @Word32 mipIx
                , " arr:" <> display @Word32 arrayIx
                , " fac:" <> display @Word32 faceIx
                , " slc:" <> display @Word32 sliceIx
                , " ]"
                ]
            let blockOffset = offset + faceIx * imageSize
            let sectionPtr = Foreign.plusPtr (VMA.mappedData stageInfo) (fromIntegral blockOffset)
            logDebug $ mconcat
              [ indices
              , " base offset = "
              , display offset
              , " image offset = "
              , display $ faceIx * imageSize
              , " image size = "
              , display imageSize
              ]
            liftIO $ unsafeUseAsCStringLen block \(pixelsPtr, pixelBytes) -> do
              if pixelBytes /= fromIntegral imageSize then
                error "assert: MipLevel.imageSize matches block.pixelBytes"
              else
                -- traceShowM (sectionPtr, pixelBytes)
                Foreign.copyBytes sectionPtr (Foreign.castPtr pixelsPtr) pixelBytes

    VMA.flushAllocation vma stage 0 Vk.WHOLE_SIZE

    -- XXX: copying to image while the staging buffer is still alive
    Image.copyBufferToImage
      context
      pool
      staging
      image
      extent
      offsets
      numLayers

  -- XXX: staging buffer is gone

  Image.transitionLayout
    context
    pool
    image
    mipLevels
    numLayers -- XXX: arrayLayers is always 0 for now
    format
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL

  imageView <- Texture.createImageView
    context
    image
    format
    mipLevels
    numLayers

  let
    extent2d =
      let
        Vk.Extent3D{..} = extent
      in
        Vk.Extent2D{..}

    allocatedImage = AllocatedImage
      { aiAllocation = allocation
      , aiExtent     = extent2d
      , aiFormat     = format
      , aiImage      = image
      , aiImageView  = imageView
      , aiImageRange = Image.subresource Vk.IMAGE_ASPECT_COLOR_BIT mipLevels numLayers
      }
  pure Texture
    { tFormat         = format
    , tMipLevels      = mipLevels
    , tLayers         = numLayers
    , tAllocatedImage = allocatedImage
    }