{-# LANGUAGE AllowAmbiguousTypes #-}

module Mischief.Render.Material.Pipeline where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State (MonadState (get), put)
import Data.Data
import Data.Foldable
import Data.Kind (Type)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Foreign (Ptr, nullPtr, with)
import Foreign.C.ConstPtr
import GHC.TypeLits
import Language.Haskell.TH (Extension (AllowAmbiguousTypes))
import Mischief.Assets (Asset (..))
import Mischief.ECS
import Mischief.Math qualified as Math
import Mischief.Render.Core
import Mischief.Render.Material.Definition
import Mischief.Render.Texture
import Mischief.Render.TextureFormat
import Mischief.Shader hiding (FragmentOutput, VertexOutput)
import Mischief.Shader.Generation hiding (FragmentOutput)
import Mischief.Shader.Internal.Bindings (Bindable, ErasedBinding (..), ErasedBindingValue, createBindGroup, createBindLayout, dummyBinding)
import Mischief.Shader.Internal.Params
import Mischief.Shader.Internal.State
import Mischief.Shader.Types qualified as ST
import Mischief.WGPU
import Mischief.WGPU.Framework (loadShaderFromBytes)
import Mischief.WGPU.Opaque
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General

data PipelineKey = PipelineKey
  { PipelineKey -> TypeRep
materialType :: TypeRep,
    PipelineKey -> WGPUTextureFormat
format :: WGPUTextureFormat
  }
  deriving (PipelineKey -> PipelineKey -> Bool
(PipelineKey -> PipelineKey -> Bool)
-> (PipelineKey -> PipelineKey -> Bool) -> Eq PipelineKey
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PipelineKey -> PipelineKey -> Bool
== :: PipelineKey -> PipelineKey -> Bool
$c/= :: PipelineKey -> PipelineKey -> Bool
/= :: PipelineKey -> PipelineKey -> Bool
Eq, Eq PipelineKey
Eq PipelineKey =>
(PipelineKey -> PipelineKey -> Ordering)
-> (PipelineKey -> PipelineKey -> Bool)
-> (PipelineKey -> PipelineKey -> Bool)
-> (PipelineKey -> PipelineKey -> Bool)
-> (PipelineKey -> PipelineKey -> Bool)
-> (PipelineKey -> PipelineKey -> PipelineKey)
-> (PipelineKey -> PipelineKey -> PipelineKey)
-> Ord PipelineKey
PipelineKey -> PipelineKey -> Bool
PipelineKey -> PipelineKey -> Ordering
PipelineKey -> PipelineKey -> PipelineKey
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: PipelineKey -> PipelineKey -> Ordering
compare :: PipelineKey -> PipelineKey -> Ordering
$c< :: PipelineKey -> PipelineKey -> Bool
< :: PipelineKey -> PipelineKey -> Bool
$c<= :: PipelineKey -> PipelineKey -> Bool
<= :: PipelineKey -> PipelineKey -> Bool
$c> :: PipelineKey -> PipelineKey -> Bool
> :: PipelineKey -> PipelineKey -> Bool
$c>= :: PipelineKey -> PipelineKey -> Bool
>= :: PipelineKey -> PipelineKey -> Bool
$cmax :: PipelineKey -> PipelineKey -> PipelineKey
max :: PipelineKey -> PipelineKey -> PipelineKey
$cmin :: PipelineKey -> PipelineKey -> PipelineKey
min :: PipelineKey -> PipelineKey -> PipelineKey
Ord)

newtype PipelineCache = PipelineCache (Map PipelineKey RenderPipeline) deriving anyclass (Typeable PipelineCache
[HookRel PipelineCache]
[Hook PipelineCache]
Set DefaultComponentType
IsExclusive (IsExclusiveRel PipelineCache)
(Typeable PipelineCache,
 IsExclusive (IsExclusiveRel PipelineCache)) =>
Set DefaultComponentType
-> [Hook PipelineCache]
-> [Hook PipelineCache]
-> [Hook PipelineCache]
-> [HookRel PipelineCache]
-> [HookRel PipelineCache]
-> [HookRel PipelineCache]
-> Component PipelineCache
forall c.
(Typeable c, IsExclusive (IsExclusiveRel c)) =>
Set DefaultComponentType
-> [Hook c]
-> [Hook c]
-> [Hook c]
-> [HookRel c]
-> [HookRel c]
-> [HookRel c]
-> Component c
$crequired :: Set DefaultComponentType
required :: Set DefaultComponentType
$conAdd :: [Hook PipelineCache]
onAdd :: [Hook PipelineCache]
$conSet :: [Hook PipelineCache]
onSet :: [Hook PipelineCache]
$conRemove :: [Hook PipelineCache]
onRemove :: [Hook PipelineCache]
$conAddRel :: [HookRel PipelineCache]
onAddRel :: [HookRel PipelineCache]
$conSetRel :: [HookRel PipelineCache]
onSetRel :: [HookRel PipelineCache]
$conRemoveRel :: [HookRel PipelineCache]
onRemoveRel :: [HookRel PipelineCache]
Component)

qgetPipeline :: forall mat. (Material mat) => TextureFormat -> Query System RenderPipeline
qgetPipeline :: forall mat.
Material mat =>
TextureFormat -> Query System RenderPipeline
qgetPipeline TextureFormat
format' = do
  let format :: WGPUTextureFormat
format = TextureFormat -> WGPUTextureFormat
toWGPUFormat TextureFormat
format'
  device <- Query System RenderDevice
forall c (m :: * -> *) w.
(MonadSystem w m, Component c) =>
Query m c
qres
  qlift $ do
    PipelineCache cache <- resOrInsert $ PipelineCache Map.empty

    let key =
          PipelineKey
            { materialType :: TypeRep
materialType = Proxy mat -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy mat -> TypeRep) -> Proxy mat -> TypeRep
forall a b. (a -> b) -> a -> b
$ forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @mat,
              WGPUTextureFormat
format :: WGPUTextureFormat
format :: WGPUTextureFormat
format
            }

    case Map.lookup key cache of
      Just RenderPipeline
pipeline -> RenderPipeline -> System RenderPipeline
forall a. a -> System a
forall (f :: * -> *) a. Applicative f => a -> f a
pure RenderPipeline
pipeline
      Maybe RenderPipeline
Nothing -> IO RenderPipeline -> System RenderPipeline
forall a. IO a -> System a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO RenderPipeline -> System RenderPipeline)
-> IO RenderPipeline -> System RenderPipeline
forall a b. (a -> b) -> a -> b
$ forall mat.
Material mat =>
RenderDevice -> WGPUTextureFormat -> IO RenderPipeline
createPipeline @mat RenderDevice
device WGPUTextureFormat
format

createPipeline :: forall mat. (Material mat) => RenderDevice -> WGPUTextureFormat -> IO RenderPipeline
createPipeline :: forall mat.
Material mat =>
RenderDevice -> WGPUTextureFormat -> IO RenderPipeline
createPipeline (RenderDevice Ptr WGPUDevice
device) WGPUTextureFormat
format = do
  BindLayout bindLayout <- [ErasedBinding] -> RenderDevice -> IO BindLayout
createBindLayout (Set ErasedBinding -> [ErasedBinding]
forall a. Set a -> [a]
Set.toList (Set ErasedBinding -> [ErasedBinding])
-> Set ErasedBinding -> [ErasedBinding]
forall a b. (a -> b) -> a -> b
$ forall m. Material m => Set ErasedBinding
collectBindings @mat) (Ptr WGPUDevice -> RenderDevice
RenderDevice Ptr WGPUDevice
device)

  pipelineLayoutDesc <- with bindLayout $ \Ptr (Ptr WGPUBindGroupLayout)
bindGroupLayout -> do
    WGPUPipelineLayoutDescriptor -> IO WGPUPipelineLayoutDescriptor
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (WGPUPipelineLayoutDescriptor -> IO WGPUPipelineLayoutDescriptor)
-> WGPUPipelineLayoutDescriptor -> IO WGPUPipelineLayoutDescriptor
forall a b. (a -> b) -> a -> b
$
      WGPUPipelineLayoutDescriptor
        { nextInChain :: Ptr WGPUChainedStruct
nextInChain = Ptr WGPUChainedStruct
forall a. Ptr a
nullPtr,
          label :: WGPUStringView
label = WGPUStringView {_data :: ConstPtr CChar
_data = Ptr CChar -> ConstPtr CChar
forall a. Ptr a -> ConstPtr a
ConstPtr Ptr CChar
forall a. Ptr a
nullPtr, length :: Int
length = Int
0},
          bindGroupLayoutCount :: CSize
bindGroupLayoutCount = CSize
1,
          bindGroupLayouts :: ConstPtr (Ptr WGPUBindGroupLayout)
bindGroupLayouts = ConstPtr {unConstPtr :: Ptr (Ptr WGPUBindGroupLayout)
unConstPtr = Ptr (Ptr WGPUBindGroupLayout)
bindGroupLayout},
          immediateSize :: CUInt
immediateSize = CUInt
0
        }

  pipelineLayout <- with pipelineLayoutDesc $ \Ptr WGPUPipelineLayoutDescriptor
pipelineLayoutDesc -> do
    Ptr WGPUDevice
-> ConstPtr WGPUPipelineLayoutDescriptor
-> IO (Ptr WGPUPipelineLayout)
wgpuDeviceCreatePipelineLayout Ptr WGPUDevice
device (ConstPtr {unConstPtr :: Ptr WGPUPipelineLayoutDescriptor
unConstPtr = Ptr WGPUPipelineLayoutDescriptor
pipelineLayoutDesc})

  -- putStrLn $ "------------------------------------------\n" ++ T.unpack (genShader @mat @(VO mat) @FragmentOutput "fragment" (collectBindings @mat) (convertFrag $ fragment @mat))

  -- putStrLn $ "------------------------------------------\n" ++ T.unpack (genShader @mat @(VI mat) @(VO mat) "vertex" (Set.toList $ collectBindings @mat) (vertex @mat))

  vertShader <- loadShaderFromBytes device (T.encodeUtf8 $ genShader @mat @(VI mat) @(VO mat) "vertex" (Set.toList $ collectBindings @mat) (vertex @mat))
  fragShader <- loadShaderFromBytes device (T.encodeUtf8 $ genShader @mat @(VO mat) @FragmentOutput "fragment" (Set.toList $ collectBindings @mat) (convertFrag $ fragment @mat))

  pipeline <- withWGPUString "main" $ \WGPUStringView
vertexEntry -> do
    String
-> (WGPUStringView -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a. String -> (WGPUStringView -> IO a) -> IO a
withWGPUString String
"main" ((WGPUStringView -> IO (Ptr WGPURenderPipeline))
 -> IO (Ptr WGPURenderPipeline))
-> (WGPUStringView -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. (a -> b) -> a -> b
$ \WGPUStringView
fragmentEntry -> do
      String
-> (WGPUStringView -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a. String -> (WGPUStringView -> IO a) -> IO a
withWGPUString String
"render pipeline" ((WGPUStringView -> IO (Ptr WGPURenderPipeline))
 -> IO (Ptr WGPURenderPipeline))
-> (WGPUStringView -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. (a -> b) -> a -> b
$ \WGPUStringView
label -> do
        let multisample :: WGPUMultisampleState
multisample =
              WGPUMultisampleState
                { nextInChain :: Ptr WGPUChainedStruct
nextInChain = Ptr WGPUChainedStruct
forall a. Ptr a
nullPtr,
                  count :: Word32
count = Word32
1,
                  mask :: Word32
mask = Word32
0xFFFFFFFF,
                  alphaToCoverageEnabled :: WGPUBool
alphaToCoverageEnabled = WGPUBool
wgpuFalse
                }

        let primitive :: WGPUPrimitiveState
primitive =
              WGPUPrimitiveState
newWGPUPrimitiveState
                { topology = wGPUPrimitiveTopology_TriangleList
                }
        WGPUBlendState
-> (Ptr WGPUBlendState -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. Storable a => a -> (Ptr a -> IO b) -> IO b
with WGPUBlendState
blendState ((Ptr WGPUBlendState -> IO (Ptr WGPURenderPipeline))
 -> IO (Ptr WGPURenderPipeline))
-> (Ptr WGPUBlendState -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. (a -> b) -> a -> b
$ \Ptr WGPUBlendState
blend -> do
          let target :: WGPUColorTargetState
target =
                WGPUColorTargetState
                  { nextInChain :: Ptr WGPUChainedStruct
nextInChain = Ptr WGPUChainedStruct
forall a. Ptr a
nullPtr,
                    WGPUTextureFormat
format :: WGPUTextureFormat
format :: WGPUTextureFormat
format,
                    blend :: ConstPtr WGPUBlendState
blend = Ptr WGPUBlendState -> ConstPtr WGPUBlendState
forall a. Ptr a -> ConstPtr a
ConstPtr Ptr WGPUBlendState
blend,
                    writeMask :: WGPUColorWriteMask
writeMask = WGPUColorWriteMask
wGPUColorWriteMask_All
                  }

          let vertex :: WGPUVertexState
vertex =
                WGPUVertexState
newVertexState
                  { entryPoint = vertexEntry,
                    _module = vertShader,
                    _WGPUVertexState = Proxy
                  }

          WGPUColorTargetState
-> (Ptr WGPUColorTargetState -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. Storable a => a -> (Ptr a -> IO b) -> IO b
with WGPUColorTargetState
target ((Ptr WGPUColorTargetState -> IO (Ptr WGPURenderPipeline))
 -> IO (Ptr WGPURenderPipeline))
-> (Ptr WGPUColorTargetState -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. (a -> b) -> a -> b
$ \Ptr WGPUColorTargetState
target -> do
            let fragment :: WGPUFragmentState
fragment =
                  WGPUFragmentState
newFragmentState
                    { entryPoint = fragmentEntry,
                      _module = fragShader,
                      _WGPUFragmentState = Proxy,
                      targetCount = 1,
                      targets = ConstPtr target
                    }

            WGPUFragmentState
-> (Ptr WGPUFragmentState -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. Storable a => a -> (Ptr a -> IO b) -> IO b
with WGPUFragmentState
fragment ((Ptr WGPUFragmentState -> IO (Ptr WGPURenderPipeline))
 -> IO (Ptr WGPURenderPipeline))
-> (Ptr WGPUFragmentState -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. (a -> b) -> a -> b
$ \Ptr WGPUFragmentState
fragment -> do
              let pipelineDesc :: WGPURenderPipelineDescriptor
pipelineDesc =
                    WGPURenderPipelineDescriptor
                      { nextInChain :: Ptr WGPUChainedStruct
nextInChain = Ptr WGPUChainedStruct
forall a. Ptr a
nullPtr,
                        WGPUStringView
label :: WGPUStringView
label :: WGPUStringView
label,
                        layout :: Ptr WGPUPipelineLayout
layout = Ptr WGPUPipelineLayout
pipelineLayout,
                        WGPUVertexState
vertex :: WGPUVertexState
vertex :: WGPUVertexState
vertex,
                        WGPUMultisampleState
multisample :: WGPUMultisampleState
multisample :: WGPUMultisampleState
multisample,
                        WGPUPrimitiveState
primitive :: WGPUPrimitiveState
primitive :: WGPUPrimitiveState
primitive,
                        fragment :: ConstPtr WGPUFragmentState
fragment = Ptr WGPUFragmentState -> ConstPtr WGPUFragmentState
forall a. Ptr a -> ConstPtr a
ConstPtr Ptr WGPUFragmentState
fragment,
                        depthStencil :: ConstPtr ()
depthStencil = Ptr () -> ConstPtr ()
forall a. Ptr a -> ConstPtr a
ConstPtr Ptr ()
forall a. Ptr a
nullPtr
                      }

              WGPURenderPipelineDescriptor
-> (Ptr WGPURenderPipelineDescriptor
    -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. Storable a => a -> (Ptr a -> IO b) -> IO b
with WGPURenderPipelineDescriptor
pipelineDesc ((Ptr WGPURenderPipelineDescriptor -> IO (Ptr WGPURenderPipeline))
 -> IO (Ptr WGPURenderPipeline))
-> (Ptr WGPURenderPipelineDescriptor
    -> IO (Ptr WGPURenderPipeline))
-> IO (Ptr WGPURenderPipeline)
forall a b. (a -> b) -> a -> b
$ \Ptr WGPURenderPipelineDescriptor
desc -> Ptr WGPUDevice
-> Ptr WGPURenderPipelineDescriptor -> IO (Ptr WGPURenderPipeline)
wgpuDeviceCreateRenderPipeline Ptr WGPUDevice
device Ptr WGPURenderPipelineDescriptor
desc

  pure $ RenderPipeline pipeline