Plugin.hs

module Mischief.Render.Plugin where

import Control.Monad
import Control.Monad.IO.Class
import Data.Data (Proxy (..))
import Data.Default
import Data.Foldable (for_)
import Data.Text qualified as T
import Data.Traversable
import Foreign (Storable (peek), free, malloc, nullPtr, with)
import Foreign.C.ConstPtr (ConstPtr (..))
import GHC.Generics
import Linear (V3 (..))
import Mischief.Assets.Asset (store)
import Mischief.ECS hiding (get, set)
import Mischief.ECS.Relationships.Order (Before (..))
import Mischief.ECS.Systems qualified as Systems
import Mischief.ECS.World.Query (queryOne)
import Mischief.Materials qualified as Mat
import Mischief.Materials.Sprite (SpriteMaterial)
import Mischief.Math (TransformPlugin)
import Mischief.Render.Bind
import Mischief.Render.Buffer
import Mischief.Render.Camera
import Mischief.Render.Core
import Mischief.Render.Material qualified as Material
import Mischief.Render.Material.Definition
import Mischief.Render.Material.Pipeline (qgetPipeline)
import Mischief.Render.Sprite (SpritePlugin (..), fetchSpriteBindings)
import Mischief.Render.Texture
import Mischief.Render.Texture qualified as Texture
import Mischief.Render.Time (TimeMaterialPlugin)
import Mischief.SDL (SDLPlugin (..))
import Mischief.SDL.Window
import Mischief.Shader (Bindable, Binding, BuiltIn, GPU, Location, Shader, ShaderParam, sample, var)
import Mischief.Shader.Internal.Bindings hiding (Bindings)
import Mischief.Shader.Types qualified as ST
import Mischief.WGPU
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General

data RenderPlugin = RenderPlugin deriving (Eq)

instance Plugin RenderPlugin where
  init = do
    first <- scheduleEntity @RenderFirst
    prepare <- scheduleEntity @RenderPrepare
    update <- scheduleEntity @RenderUpdate
    transparent <- scheduleEntity @RenderTransparent
    postprocess <- scheduleEntity @RenderPostprocess
    last <- scheduleEntity @RenderLast

    update' <- scheduleEntity @PostUpdate

    insert (Rel Before first) update'
    insert (Rel Before prepare) first
    insert (Rel Before update) prepare
    insert (Rel Before transparent) update
    insert (Rel Before postprocess) transparent
    insert (Rel Before last) postprocess

    for_ [first, update, prepare, transparent, postprocess, last] $ insert UpdateSchedule

    insertRes =<< liftIO (RenderInstance <$> wgpuCreateInstance)
    void $ spawn $ Observer onAddWindow
    void $ spawn $ Observer onSetWindowSize

    systems renderCameras
      & schedule @RenderLast

    systems clearCameras
      & schedule @RenderFirst

  deps = [dep @SDLPlugin, dep @SpritePlugin, dep @CameraPlugin, dep @TimeMaterialPlugin, dep @TransformPlugin]

-- TODO: refactor this to use the query monad
renderCameras :: System ()
renderCameras = do
  resources <- getRenderingResources
  case resources of
    Nothing -> pure ()
    Just (adapter, device, queue) -> do
      cameras <- query [q|CameraTexture, OutputTo -> (RenderSurface)|]
      for_ cameras $ \(CameraTexture SwappableTexture {currentTexture = texture}, From _ surface) -> do
        format <- getFormat surface adapter
        withSurfaceTexture surface (def {Texture.format}) $ \output -> do
          sampler <- newSampler device
          tex <- liftIO $ wgpuTextureCreateView texture (ConstPtr nullPtr)
          let command = Draw [ErasedBindingValue Bindings {tex = TextureView tex, sampler}] 3
          Just pipeline <- queryOne $ qgetPipeline @DummyFullscreenMaterial format
          render @DummyFullscreenMaterial device queue pipeline def output [command]
          presentSurface surface

clearCameras :: System ()
clearCameras = query_ $ do
  device <- qres
  queue <- qres
  (CameraTexture SwappableTexture {currentTexture = texture, desc}) <- [q|CameraTexture|]
  pipeline <- qgetPipeline @DummyFullscreenMaterial desc.format
  qlift $ render @ClearMaterial device queue pipeline def {clear = Clear $ V3 0.184313725 0.211764706 0.250980392} (Texture texture desc) []

withSurfaceTexture :: RenderSurface -> TextureDescriptor -> (Texture -> System a) -> System ()
withSurfaceTexture (RenderSurface surface) desc f = do
  surfaceTextureBox <- liftIO malloc
  liftIO $ wgpuSurfaceGetCurrentTexture surface surfaceTextureBox
  surfaceTexture <- liftIO $ peek surfaceTextureBox

  unless (surfaceTexture.texture == nullPtr) $ void $ f (Texture {texture = surfaceTexture.texture, desc})

  liftIO $ free surfaceTextureBox

presentSurface :: RenderSurface -> System ()
presentSurface (RenderSurface surface) = liftIO $ wgpuSurfacePresent surface

data Bindings f = Bindings
  { tex :: Binding f Texture,
    sampler :: Binding f Sampler
  }
  deriving (Generic, Bindable)

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

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

-- | Material for rendering something to the whole screen.
data DummyFullscreenMaterial

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

  vertex :: VertexInput GPU -> Shader DummyFullscreenMaterial (VertexOutput GPU)
  vertex = liftExt . vertexShader

  fragment :: VertexOutput GPU -> Shader DummyFullscreenMaterial (Mat.Vec 4 Mat.Float)
  fragment input = do
    b <- Mat.bindings
    pure $ Mat.sample b.tex b.sampler input.uv

data ClearMaterial

data NoBindings f = NoBindings deriving (Generic, Bindable)

instance Material ClearMaterial where
  type VI ClearMaterial = VertexInput
  type VO ClearMaterial = VertexOutput
  type B ClearMaterial = NoBindings
  vertex = liftExt . vertexShader

  fragment _ = pure $ Mat.vec4 0 0 0 0

vertexShader :: VertexInput GPU -> Shader NoMaterial (VertexOutput GPU)
vertexShader (VertexInput index) = do
  positions <- Mat.var $ Mat.array @(Mat.Vec 2 Mat.Float) [Mat.vec2 (-1) (-1), Mat.vec2 3 (-1), Mat.vec2 (-1) 3]
  let pos = positions Mat.!! index
  pure $
    VertexOutput
      { pos = Mat.vec4 pos.x pos.y 0 1,
        uv = Mat.vec2 (pos.x Mat.* 0.5 + 0.5) (1.0 Mat.- (pos.y Mat.* 0.5 Mat.+ 0.5))
      }