Bindings.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -Wno-partial-fields #-}

module Mischief.Shader.Internal.Bindings where

import Data.Bifunctor qualified
import Data.Data
import Data.Default (Default (def))
import Data.Kind
import Data.Primitive.Ptr
import Data.Singletons (SingI)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Vector.Storable qualified as VS
import Foreign (Bits ((.|.)), with)
import Foreign.C.ConstPtr
import GHC.Generics (Generic (Rep, from), K1 (K1), M1 (M1), U1 (U1), V1 (..), to, (:*:) ((:*:)))
import GHC.TypeLits
import Mischief.ECS (System)
import Mischief.Render.Buffer (Buffer (..), BufferDescriptor (..))
import Mischief.Render.Core
import Mischief.Render.Texture
import Mischief.Shader.Internal.Buffers qualified as Buffers
import Mischief.Shader.Internal.Singletons
import Mischief.Shader.Internal.State
import Mischief.Shader.Types hiding ((<>))
import Mischief.WGPU
import Mischief.WGPU.Opaque (WGPUBindGroup, WGPUBindGroupLayout)
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General

-- data Binding (index :: Nat) a where
--   Binding :: forall b index. Expr (AssociatedExpr b) -> Binding index b
--   BindEntry :: WGPUBindGroupEntry -> Binding index b

-- data Uniform (index :: Nat) a where
--   Uniform :: forall b index. WGPUBindGroupEntry -> Uniform index b

newtype UniformBinding a = UniformBinding a

newtype NormalBinding a = NormalBinding a

type family Binding f a where
  Binding Internal b = NormalBinding b
  Binding CPU Sampler = Sampler
  Binding GPU Sampler = Expr TSampler
  Binding CPU Texture = TextureView
  Binding GPU Texture = Expr TTexture

type family Uniform f a where
  Uniform Internal b = UniformBinding b
  Uniform CPU (a b) = Buffer a
  Uniform GPU (Expr a) = Expr a
  Uniform GPU (a GPU) = a GPU

-- Uniform CPU

-- type family ToGpu a where
--   T

type family AssociatedExpr a where
  AssociatedExpr (Expr a) = a
  AssociatedExpr Texture = TTexture
  AssociatedExpr Sampler = TSampler

-- instance (KnownNat n, SingI (AssociatedExpr a)) => Default (Binding n a) where
--   def = Binding $ BindingVar (natVal (Proxy @n))

data BindingData
  = NormalData
      { bType :: Text,
        index :: Integer
      }
  | UniformData
      { structName :: Text,
        structFields :: [(Text, Text)],
        index :: Integer
      }
  deriving (Show)

newtype Bindings = Bindings [BindingData] deriving newtype (Semigroup, Show)

class (Typeable a) => Bindable a where
  collectBindings :: Nat -> Proxy (a Internal) -> (Bindings, Nat)
  default collectBindings :: (GCollectBindings (Rep (a Internal))) => Nat -> Proxy (a Internal) -> (Bindings, Nat)
  collectBindings n _ = gCollectBindings (Proxy @(Rep (a Internal))) n

  collectLayouts :: Nat -> Proxy (a Internal) -> ([WGPUBindGroupLayoutEntry], Nat)
  default collectLayouts :: (GCollectBindings (Rep (a Internal))) => Nat -> Proxy (a Internal) -> ([WGPUBindGroupLayoutEntry], Nat)
  collectLayouts n _ = gCollectLayouts (Proxy @(Rep (a Internal))) n

  collectEntries :: Nat -> a CPU -> ([WGPUBindGroupEntry], Nat)
  default collectEntries :: (GCollectEntries (Rep (a CPU)) (Rep (a Internal)), Generic (a CPU)) => Nat -> a CPU -> ([WGPUBindGroupEntry], Nat)
  collectEntries n a = gCollectEntries @(Rep (a CPU)) @(Rep (a Internal)) (from a) n

  dummyBinding :: Nat -> (a GPU, Nat)
  default dummyBinding :: (Generic (a GPU), GDummyBinding (Rep (a Internal)) (Rep (a GPU))) => Nat -> (a GPU, Nat)
  dummyBinding n = Data.Bifunctor.first to $ gDummyBinding @(Rep (a Internal)) @(Rep (a GPU)) n

data ErasedBinding where
  ErasedBinding :: forall b (a :: Type) m. (Typeable m, Bindable b, Typeable b, Typeable a) => {materialType :: m, material :: Proxy a, binding :: Proxy b} -> ErasedBinding

instance Eq ErasedBinding where
  ErasedBinding _ a _ == ErasedBinding _ b _ = typeRep a == typeRep b

instance Ord ErasedBinding where
  compare (ErasedBinding _ a _) (ErasedBinding _ b _) = compare (typeRep a) (typeRep b)

data ErasedBindingValue where
  ErasedBindingValue :: (Bindable b, Typeable b) => b CPU -> ErasedBindingValue

class CollectBindings a where
  collectBinding :: Proxy a -> Nat -> BindingData
  collectLayout :: Proxy a -> Nat -> WGPUBindGroupLayoutEntry

class GCollectBindings f where
  gCollectBindings :: Proxy f -> Nat -> (Bindings, Nat)
  gCollectLayouts :: Proxy f -> Nat -> ([WGPUBindGroupLayoutEntry], Nat)

instance GCollectBindings V1 where
  gCollectBindings _ n = (Bindings [], n)
  gCollectLayouts _ n = ([], n)

instance GCollectBindings U1 where
  gCollectBindings _ n = (Bindings [], n)
  gCollectLayouts _ n = ([], n)

instance (GCollectBindings f, GCollectBindings g) => GCollectBindings (f :*: g) where
  gCollectBindings _ n = do
    let (a, n') = gCollectBindings (Proxy @f) n
    let (b, n'') = gCollectBindings (Proxy @g) n'
    (a <> b, n'')
  gCollectLayouts _ n = do
    let (a, n') = gCollectLayouts (Proxy @f) n
    let (b, n'') = gCollectLayouts (Proxy @g) n'
    (a <> b, n'')

instance (CollectBindings c) => GCollectBindings (K1 i c) where
  gCollectBindings _ n = (Bindings [collectBinding (Proxy @c) n], n + 1)
  gCollectLayouts _ n = ([collectLayout (Proxy @c) n], n + 1)

instance (GCollectBindings f) => GCollectBindings (M1 i t f) where
  gCollectBindings _ = gCollectBindings (Proxy @f)
  gCollectLayouts _ = gCollectLayouts (Proxy @f)

instance CollectBindings (NormalBinding Sampler) where
  collectBinding _ n = NormalData {bType = T.pack "sampler", index = fromIntegral n}
  collectLayout _ n =
    WGPUBindGroupLayoutEntry
      { nextInChain = nullPtr,
        binding = fromIntegral n,
        visibility = wGPUShaderStage_Fragment .|. wGPUShaderStage_Vertex,
        bindingArraySize = 0,
        buffer = unusedBufferLayout,
        storageTexture = unusedStorageTextureLayout,
        texture = unusedTextureLayout,
        sampler =
          WGPUSamplerBindingLayout
            { _type = wGPUSamplerBindingType_Filtering,
              nextInChain = nullPtr
            }
      }

instance CollectBindings (NormalBinding Texture) where
  collectBinding _ n = NormalData {bType = T.pack "texture_2d<f32>", index = fromIntegral n}
  collectLayout _ n =
    WGPUBindGroupLayoutEntry
      { nextInChain = nullPtr,
        binding = fromIntegral n,
        visibility = wGPUShaderStage_Fragment .|. wGPUShaderStage_Vertex,
        bindingArraySize = 0,
        buffer = unusedBufferLayout,
        sampler = unusedSamplerLayout,
        storageTexture = unusedStorageTextureLayout,
        texture =
          WGPUTextureBindingLayout
            { sampleType = wGPUTextureSampleType_Float,
              multisampled = wgpuFalse,
              viewDimension = wGPUTextureViewDimension_2D,
              nextInChain = nullPtr
            }
      }

instance (Buffers.Bufferable a, Typeable a) => CollectBindings (UniformBinding (a Internal)) where
  collectBinding _ n =
    UniformData
      { structName = T.pack $ show $ typeRep (Proxy @a),
        structFields = Buffers.getFields @a,
        index = fromIntegral n
      }

  collectLayout _ n =
    WGPUBindGroupLayoutEntry
      { nextInChain = nullPtr,
        binding = fromIntegral n,
        visibility = wGPUShaderStage_Fragment .|. wGPUShaderStage_Vertex,
        bindingArraySize = 0,
        buffer =
          WGPUBufferBindingLayout
            { nextInChain = nullPtr,
              _type = wGPUBufferBindingType_Uniform,
              hasDynamicOffset = wgpuFalse,
              minBindingSize = fromIntegral $ Buffers.getSize @a
            },
        sampler = unusedSamplerLayout,
        storageTexture = unusedStorageTextureLayout,
        texture = unusedTextureLayout
      }

class CollectEntry a n where
  collectEntry :: a -> Nat -> WGPUBindGroupEntry

class GCollectEntries f g where
  gCollectEntries :: f p -> Nat -> ([WGPUBindGroupEntry], Nat)

instance GCollectEntries V1 V1 where
  gCollectEntries _ n = ([], n)

instance GCollectEntries U1 U1 where
  gCollectEntries _ n = ([], n)

instance (GCollectEntries f1 g1, GCollectEntries f2 g2) => GCollectEntries (f1 :*: f2) (g1 :*: g2) where
  gCollectEntries (f :*: g) n = do
    let (a, n') = gCollectEntries @f1 @g1 f n
    let (b, n'') = gCollectEntries @f2 @g2 g n'
    (a <> b, n'')

instance (CollectEntry c n) => GCollectEntries (K1 i c) (K1 i n) where
  gCollectEntries (K1 a) n = ([collectEntry @c @n a n], n + 1)

instance (GCollectEntries f n) => GCollectEntries (M1 i t f) (M1 i t n) where
  gCollectEntries (M1 a) = gCollectEntries @f @n a

instance CollectEntry TextureView (NormalBinding b) where
  collectEntry (TextureView textureView) n =
    WGPUBindGroupEntry
      { binding = fromIntegral n,
        textureView,
        size = 0,
        offset = 0,
        sampler = nullPtr,
        buffer = nullPtr,
        nextInChain = nullPtr
      }

instance CollectEntry Sampler (NormalBinding b) where
  collectEntry (Sampler sampler) n =
    WGPUBindGroupEntry
      { binding = fromIntegral n,
        textureView = nullPtr,
        size = 0,
        offset = 0,
        sampler,
        buffer = nullPtr,
        nextInChain = nullPtr
      }

instance CollectEntry (Buffer a) (UniformBinding b) where
  collectEntry Buffer {buffer, desc = BufferDescriptor {size}} n =
    WGPUBindGroupEntry
      { binding = fromIntegral n,
        textureView = nullPtr,
        size = fromIntegral size,
        offset = 0,
        sampler = nullPtr,
        buffer,
        nextInChain = nullPtr
      }

class DummyBinding a b where
  dummyBinding' :: Nat -> b

class GDummyBinding f g where
  gDummyBinding :: Nat -> (g p, Nat)

instance GDummyBinding V1 V1 where
  gDummyBinding = undefined

instance GDummyBinding U1 U1 where
  gDummyBinding n = (U1, n)

instance (GDummyBinding f1 g1, GDummyBinding f2 g2) => GDummyBinding (f1 :*: f2) (g1 :*: g2) where
  gDummyBinding n = do
    let (a, n') = gDummyBinding @f1 @g1 n
    let (b, n'') = gDummyBinding @f2 @g2 n'
    (a :*: b, n'')

instance (DummyBinding a b) => GDummyBinding (K1 i a) (K1 i b) where
  gDummyBinding n = let a = dummyBinding' @a @b n in (K1 a, n + 1)

instance (GDummyBinding a b) => GDummyBinding (M1 i t a) (M1 i t b) where
  gDummyBinding n = let (a, n') = gDummyBinding @a @b n in (M1 a, n')

instance (SingI b) => DummyBinding (NormalBinding a) (Expr b) where
  dummyBinding' n = BindingVar (fromIntegral n)

instance (Buffers.Bufferable a) => DummyBinding (UniformBinding (a Internal)) (a GPU) where
  dummyBinding' n = Buffers.dummyBuffer @a (fromIntegral n)

-- instance (KnownNat n) => Bindable (UniformBinding n Sampler) where
--   collectBindings _ = Bindings [BindingData {bType = T.pack "sampler", index = natVal (Proxy @n)}]
--   collectLayouts _ =
--     [ WGPUBindGroupLayoutEntry
--         { nextInChain = nullPtr,
--           binding = fromInteger (natVal (Proxy @n)),
--           visibility = wGPUShaderStage_Fragment,
--           bindingArraySize = 0,
--           buffer = unusedBufferLayout,
--           storageTexture = unusedStorageTextureLayout,
--           texture = unusedTextureLayout,
--           sampler =
--             WGPUSamplerBindingLayout
--               { _type = wGPUSamplerBindingType_Filtering,
--                 nextInChain = nullPtr
--               }
--         }
--     ]

--   collectEntries (UniformBinding a) = [bind a]

-- -- instance (KnownNat n, SingI b, f Internal ~ UniformBinding n a', f GPU ~ Expr b) => ToDummy f where
-- --   dummyB = BindingVar (natVal (Proxy @n))

-- instance (KnownNat n) => Bindable (UniformBinding n Texture) where
--   collectBindings _ = Bindings [BindingData {bType = T.pack "texture_2d<f32>", index = natVal (Proxy @n)}]
--   collectLayouts _ =
--     [ WGPUBindGroupLayoutEntry
--         { nextInChain = nullPtr,
--           binding = fromInteger (natVal (Proxy @n)),
--           visibility = wGPUShaderStage_Fragment,
--           bindingArraySize = 0,
--           buffer = unusedBufferLayout,
--           sampler = unusedSamplerLayout,
--           storageTexture = unusedStorageTextureLayout,
--           texture =
--             WGPUTextureBindingLayout
--               { sampleType = wGPUTextureSampleType_Float,
--                 multisampled = wgpuFalse,
--                 viewDimension = wGPUTextureViewDimension_2D,
--                 nextInChain = nullPtr
--               }
--         }
--     ]

--   collectEntries (UniformBinding a) = [bind a]

-- -- dummyB = Binding $ BindingVar (natVal (Proxy @n))

-- instance Bindable ()

-- instance (KnownNat n, SingI (Primitive a)) => Bindable (UniformBinding n (Expr (Primitive a))) where
--   collectBindings _ = Bindings [BindingData {bType = typeToWGSL (undefined :: Expr (Primitive a)), index = natVal (Proxy @n)}]
--   collectLayouts _ = undefined

--   collectEntries (UniformBinding a) = [bind a]

-- -- dummyB = Binding $ BindingVar (natVal (Proxy @n))

-- instance (KnownNat n, SingI (VectorType m a)) => Bindable (UniformBinding n (Expr (VectorType m a))) where
--   collectBindings _ = Bindings [BindingData {bType = typeToWGSL (undefined :: Expr (VectorType m a)), index = natVal (Proxy @n)}]
--   collectLayouts _ = undefined

--   collectEntries (UniformBinding a) = [bind a]

-- -- dummyB = Binding $ BindingVar (natVal (Proxy @n))

-- instance (KnownNat n, SingI (ArrayType m a)) => Bindable (UniformBinding n (Expr (ArrayType m a))) where
--   collectBindings _ = Bindings [BindingData {bType = typeToWGSL (undefined :: Expr (ArrayType m a)), index = natVal (Proxy @n)}]
--   collectLayouts _ = undefined

--   collectEntries (UniformBinding a) = [bind a]

-- -- dummyB = Binding $ BindingVar (natVal (Proxy @n))

collectAllLayouts :: Nat -> [ErasedBinding] -> [WGPUBindGroupLayoutEntry]
collectAllLayouts _ [] = []
collectAllLayouts n (ErasedBinding _ (_ :: Proxy m) (_ :: Proxy b) : xs) =
  let (x, n') = collectLayouts n (Proxy @(b Internal))
   in x ++ collectAllLayouts n' xs

createBindLayout :: [ErasedBinding] -> RenderDevice -> IO BindLayout
createBindLayout b (RenderDevice device) = do
  let layouts = VS.fromList $ collectAllLayouts 0 b

  withWGPUString "bind group layout" $ \label -> do
    x <- VS.unsafeWith layouts $ \entries -> do
      let desc =
            WGPUBindGroupLayoutDescriptor
              { entryCount = fromIntegral $ VS.length layouts,
                entries = ConstPtr entries,
                label,
                nextInChain = nullPtr
              }
      with desc $ \desc -> wgpuDeviceCreateBindGroupLayout device (ConstPtr desc)

    pure (BindLayout x)

collectAllEntries :: Nat -> [ErasedBindingValue] -> [WGPUBindGroupEntry]
collectAllEntries _ [] = []
collectAllEntries n (ErasedBindingValue b : xs) =
  let (x, n') = collectEntries n b
   in x ++ collectAllEntries n' xs

createBindGroup :: RenderDevice -> BindLayout -> [ErasedBindingValue] -> IO (Ptr WGPUBindGroup)
createBindGroup (RenderDevice device) (BindLayout bindGroupLayout) b = do
  let entries = VS.fromList $ collectAllEntries 0 b

  withWGPUString "bind group" $ \label -> do
    VS.unsafeWith entries $ \entries' -> do
      let desc =
            WGPUBindGroupDescriptor
              { layout = bindGroupLayout,
                entryCount = fromIntegral $ VS.length entries,
                label = label,
                nextInChain = nullPtr,
                entries = ConstPtr entries'
              }

      with desc $ \desc -> wgpuDeviceCreateBindGroup device (ConstPtr desc)

-- class (Bindable link) => ConvertBindings link where
--   type CpuBind link
--   type GpuBind link
--   toLink :: CpuBind link -> link
--   mkGpu :: GpuBind link

-- bind :: a -> WGPUBindGroupEntry
-- bind = undefined

-- instance (KnownNat n) => ConvertBindings (UniformBinding (n :: Nat) Texture) where
--   type CpuBind (UniformBinding (n :: Nat) Texture) = TextureView
--   type GpuBind (UniformBinding (n :: Nat) Texture) = Expr TTexture
--   toLink (TextureView view) =
--     Uniform
--       WGPUBindGroupEntry
--         { binding = fromInteger (natVal (Proxy @n)),
--           textureView = view,
--           size = 0,
--           offset = 0,
--           sampler = nullPtr,
--           buffer = nullPtr,
--           nextInChain = nullPtr
--         }

--   mkGpu = BindingVar (natVal (Proxy @n))

-- instance (KnownNat n) => ConvertBindings (Uniform (n :: Nat) Sampler) where
--   type CpuBind (Uniform (n :: Nat) Sampler) = Sampler
--   type GpuBind (Uniform (n :: Nat) Sampler) = Expr TSampler
--   toLink (Sampler sampler) =
--     Uniform
--       WGPUBindGroupEntry
--         { binding = fromInteger (natVal (Proxy @n)),
--           textureView = nullPtr,
--           size = 0,
--           offset = 0,
--           sampler = sampler,
--           buffer = nullPtr,
--           nextInChain = nullPtr
--         }

--   mkGpu = BindingVar (natVal (Proxy @n))

-- class GMkDummy i o where
--   gMkDummy :: o

-- instance GMkDummy (K1 a b c) (K1 a b c) where