Sprite.hs

module Mischief.Materials.Sprite where

import Mischief.Materials
import Mischief.Materials.Time (TimeMaterial, getTime)
import Mischief.Render.Camera (CameraMatrices (CameraMatrices), Matrices (..))
import Mischief.Render.Core (DefaultSampler (..))
import Mischief.Render.Image
import Mischief.Render.Sprite.Common
import Prelude ()

data SpriteData f = SpriteData
  { coords :: Field f (Vec 3 Float),
    size :: Field f (Vec 2 Float),
    slice :: Field f (Vec 4 Float),
    flip :: Field f (Vec 2 Int),
    rotation :: Field f (Mat 3 3)
  }
  deriving stock (Generic)
  deriving anyclass (Bufferable)

newtype VertexInput f = VertexInput (BuiltIn f "vertex_index" Uint)
  deriving stock (Generic)
  deriving anyclass (ShaderParam)

data VertexOutput f = VertexOutput
  { pos :: BuiltIn f "position" (Vec 4 Float),
    uv :: Location f 0 (Vec 2 Float)
  }
  deriving (Generic, ShaderParam)

data Bindings f = Bindings
  { matrices :: Uniform f (Matrices f),
    sprite :: Uniform f (SpriteData f),
    texture :: Binding f Texture,
    sampler :: Binding f Sampler
  }
  deriving (Generic, Bindable)

data SpriteMaterial

instance Material SpriteMaterial where
  type VI SpriteMaterial = VertexInput
  type VO SpriteMaterial = VertexOutput
  type B SpriteMaterial = Bindings

  vertex :: VertexInput GPU -> Shader SpriteMaterial (VertexOutput GPU)
  vertex (VertexInput index) = do
    b <- bindings

    positions <- var $ array @(Vec2 Float) [vec2 (-1) (-1), vec2 1 (-1), vec2 1 1, vec2 (-1) (-1), vec2 1 1, vec2 (-1) 1]
    uvs <- var $ array @(Vec2 Float) [vec2 0 1, vec2 1 1, vec2 1 0, vec2 0 1, vec2 1 0, vec2 0 0]

    pos <- var $ (positions !! index) * b.sprite.size * diag 0.5
    pos <- var $ b.sprite.rotation !*^ (pos +: 0)
    pos <- var $ pos + b.sprite.coords
    pos <- var $ b.matrices.projection !*^ (b.matrices.view !*^ (pos +: 1))

    pure
      VertexOutput
        { pos = vec3 pos.x pos.y pos.z +: 1,
          uv = uvs !! index
        }

  fragment :: VertexOutput GPU -> Shader SpriteMaterial (Vec 4 Float)
  fragment VertexOutput {uv} = sampleSprite uv

  fetch =
    FetchBindings
      ( \camera -> do
          (DefaultSampler sampler) <- qres @DefaultSampler
          CameraMatrices matrices <- qget camera [qd|CameraMatrices|]
          pure (matrices, sampler)
      )
      ( \sprite -> do
          (Sprite {image}, buffer) <- qget sprite [qd|Sprite, (Buffer SpriteData)|]
          ImageTextureView texture <- qget image [qd|ImageTextureView|]
          pure (buffer, texture)
      )
      (\(matrices, sampler) (buffer, texture) -> Bindings {matrices, sampler, sprite = buffer, texture})

sampleSprite :: Vec 2 Float -> Shader SpriteMaterial (Vec 4 Float)
sampleSprite uv = do
  b <- bindings
  let uv_flipped :: Vec 2 Float = abs $ cast b.sprite.flip - uv
  let uv = uv_flipped * vec2 b.sprite.slice.z b.sprite.slice.w + b.sprite.slice.xy
  pure $ sample b.texture b.sampler uv

newtype SomeColor f = SomeColor {color :: Field f (Vec 3 Float)}
  deriving stock (Generic)
  deriving anyclass (Bufferable)

newtype CustomBindings f = CustomBindings (Uniform f (SomeColor f))
  deriving stock (Generic)
  deriving anyclass (Bindable)

data CustomMat

instance Material CustomMat where
  type VI CustomMat = VI SpriteMaterial
  type VO CustomMat = VO SpriteMaterial

  type B CustomMat = CustomBindings

  type Extends CustomMat = '[SpriteMaterial, TimeMaterial]

  vertex = liftExt . vertex @SpriteMaterial

  fragment :: VO CustomMat GPU -> Shader CustomMat (Vec 4 Float)
  fragment vo = do
    color <- liftExt $ fragment @SpriteMaterial vo
    time <- liftExt getTime
    (CustomBindings (SomeColor c)) <- bindings
    pure $ color * (c ^* cos time +: 1)

  fetch =
    FetchBindings
      { camera = const $ pure (),
        mesh = \sprite -> qget sprite [qd|(Buffer SomeColor)|],
        result = const CustomBindings
      }