Texture.hs

module Mischief.Render.Texture where

import Codec.Picture qualified as P
import Codec.Picture.Extra (scaleBilinear)
import Control.Applicative.Singletons
import Control.Monad.IO.Class
import Data.Default (Default (def))
import Data.Maybe.Singletons (SMaybe (SJust))
import Data.Singletons
import Data.Vector.Storable (Vector)
import Data.Vector.Storable qualified as VS
import Foreign (Bits ((.|.)), Ptr, Word32, Word8, castPtr, nullPtr, with)
import Foreign.C.ConstPtr (ConstPtr (..))
import GHC.TypeLits
import Mischief.Assets
import Mischief.Assets.Image
import Mischief.ECS
import Mischief.Render.Core
import Mischief.Render.TextureFormat
import Mischief.WGPU
import Mischief.WGPU.Opaque
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General

data Texture = Texture {texture :: Ptr WGPUTexture, desc :: TextureDescriptor}

data SwappableTexture = SwappableTexture {currentTexture :: Ptr WGPUTexture, otherTexture :: Ptr WGPUTexture, desc :: TextureDescriptor}

extractTexture :: SwappableTexture -> Texture
extractTexture SwappableTexture {currentTexture, desc} = Texture currentTexture desc

swapTexture :: SwappableTexture -> (SwappableTexture, Texture)
swapTexture SwappableTexture {currentTexture, otherTexture, desc} = (SwappableTexture {currentTexture = otherTexture, otherTexture = currentTexture, desc}, Texture currentTexture desc)

newtype TextureView = TextureView (Ptr WGPUTextureView)

data TextureDescriptor = TextureDescriptor
  { width :: Int,
    height :: Int,
    depth :: Int,
    label :: String,
    mipLevels :: Int,
    samples :: Int,
    format :: TextureFormat,
    usages :: [TextureUsage]
  }
  deriving (Show)

instance Default TextureDescriptor where
  def = textureDescriptor

textureDescriptor :: TextureDescriptor
textureDescriptor =
  TextureDescriptor
    { width = 1,
      height = 1,
      depth = 1,
      label = "texture",
      mipLevels = 1,
      samples = 1,
      format = RGBA8UnormSrgb,
      usages = [TextureUsageTextureBinding, TextureUsageCopyDst, TextureUsageRenderAttachment]
    }

data TextureUsage
  = TextureUsageNone
  | TextureUsageCopySrc
  | TextureUsageCopyDst
  | TextureUsageTextureBinding
  | TextureUsageStorageBinding
  | TextureUsageRenderAttachment
  | TextureUsageTransientAttachment
  deriving (Show)

usageBits :: TextureUsage -> WGPUTextureUsage
usageBits TextureUsageNone = wGPUTextureUsage_None
usageBits TextureUsageCopySrc = wGPUTextureUsage_CopySrc
usageBits TextureUsageCopyDst = wGPUTextureUsage_CopyDst
usageBits TextureUsageTextureBinding = wGPUTextureUsage_TextureBinding
usageBits TextureUsageStorageBinding = wGPUTextureUsage_StorageBinding
usageBits TextureUsageRenderAttachment = wGPUTextureUsage_RenderAttachment
usageBits TextureUsageTransientAttachment = wGPUTextureUsage_TransientAttachment

processUsages :: [TextureUsage] -> WGPUTextureUsage
processUsages = foldr ((.|.) . usageBits) (WGPUTextureUsage (WGPUFlags 0))

createSwappableTexture :: RenderDevice -> TextureDescriptor -> IO SwappableTexture
createSwappableTexture device desc = do
  a <- createTexture' device desc
  b <- createTexture' device desc
  pure $ SwappableTexture a b desc

createTexture :: RenderDevice -> TextureDescriptor -> IO Texture
createTexture device desc = (`Texture` desc) <$> createTexture' device desc

createTexture' :: RenderDevice -> TextureDescriptor -> IO (Ptr WGPUTexture)
createTexture' (RenderDevice device) descriptor = liftIO $ do
  let TextureDescriptor {width, height, depth, label, mipLevels, samples, format, usages} = descriptor
  withWGPUString label $ \label -> do
    let desc =
          WGPUTextureDescriptor
            { size = WGPUExtent3D (fromIntegral width) (fromIntegral height) (fromIntegral depth),
              label,
              mipLevelCount = fromIntegral mipLevels,
              sampleCount = fromIntegral samples,
              dimension = wGPUTextureDimension_2D,
              format = toWGPUFormat format,
              usage = processUsages usages,
              nextInChain = nullPtr,
              viewFormatCount = 0,
              viewFormats = ConstPtr nullPtr
            }

    with desc (wgpuDeviceCreateTexture device . ConstPtr)

getDimension :: (Int, Int, Int) -> WGPUTextureDimension
getDimension (_, 1, 1) = wGPUTextureDimension_1D
getDimension (_, _, 1) = wGPUTextureDimension_2D
getDimension (_, _, _) = wGPUTextureDimension_3D

uploadImage :: RenderQueue -> Image -> Texture -> IO ()
uploadImage (RenderQueue queue) (Image image) (Texture {texture, desc = TextureDescriptor {width, height}}) = do
  let (P.Image w h bytes) = scaleBilinear width height (P.convertRGBA8 image)

  let extent = WGPUExtent3D (fromIntegral w) (fromIntegral h) 1
  let copyInfo =
        WGPUTexelCopyTextureInfo
          { texture,
            mipLevel = 0,
            origin = WGPUOrigin3D 0 0 0,
            aspect = wGPUTextureAspect_All
          }

  let layout =
        WGPUTexelCopyBufferLayout
          { offset = 0,
            bytesPerRow = 4 * fromIntegral w,
            rowsPerImage = fromIntegral h
          }

  VS.unsafeWith bytes $ \pixelPtr -> do
    with extent $ \extent -> do
      with copyInfo $ \copyInfo -> do
        with layout $ \layout -> do
          wgpuQueueWriteTexture queue (ConstPtr copyInfo) (ConstPtr $ castPtr pixelPtr) (fromIntegral $ VS.length bytes) (ConstPtr layout) (ConstPtr extent)

createTextureForImage :: RenderDevice -> Image -> IO Texture
createTextureForImage device (Image image) = do
  let (P.Image w h _) = P.convertRGBA8 image
  createTexture
    device
    TextureDescriptor
      { width = w,
        height = h,
        depth = 1,
        label = "texture",
        mipLevels = 1,
        samples = 1,
        format = RGBA8UnormSrgb,
        usages = [TextureUsageTextureBinding, TextureUsageCopyDst]
      }

createView :: Texture -> System TextureView
createView Texture {texture} = liftIO $ TextureView <$> wgpuTextureCreateView texture (ConstPtr nullPtr)