Definition.hs

{-# LANGUAGE AllowAmbiguousTypes #-}

module Mischief.Render.Material.Definition 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 Mischief.Assets (Asset (..))
import Mischief.ECS
import Mischief.Math qualified as Math
import Mischief.Render.Core
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
import Unsafe.Coerce (unsafeCoerce)

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

newtype DefaultVO f = DefaultVO
  { pos :: BuiltIn f "position" (Vec 4 ST.Float)
  }
  deriving (Generic)
  deriving anyclass (ShaderParam)

class
  ( Bindable (B a),
    ShaderParam (VI a),
    ShaderParam (VO a),
    Extension (Extends a),
    Typeable a,
    Typeable (B a),
    Typeable (VI a),
    Typeable (VO a),
    Typeable (Extends a),
    Typeable (IsPhantom a)
  ) =>
  Material (a :: Type)
  where
  type VI a :: Type -> Type
  type VO a :: Type -> Type

  type B a :: Type -> Type

  type Extends a :: [Type]
  type Extends a = '[]

  type Extandable a :: Bool
  type Extandable a = True

  type IsPhantom a :: Bool
  type IsPhantom a = False

  vertex :: VI a GPU -> Shader a (VO a GPU)
  default vertex :: (IsPhantom a ~ True) => VI a GPU -> Shader a (VO a GPU)
  vertex = undefined

  fragment :: VO a GPU -> Shader a (ST.Vec 4 ST.Float)
  default fragment :: (IsPhantom a ~ True) => VO a GPU -> Shader a (ST.Vec 4 ST.Float)
  fragment = undefined

  fetch :: FetchBindings (B a)

data PhantomMaterial = PhantomMaterial deriving (Component)

data MaterialType where
  MaterialType :: forall m. (Material m) => Proxy m -> MaterialType
  deriving (Component)

data NoMaterial

-- type StoreData MaterialType = (MaterialType, )
-- storeAsset = undefined

liftExt :: forall a b c. (ExtensionOf b a) => Shader a c -> Shader b c
liftExt (Shader a) = Shader a

class Extension (a :: [Type]) where
  collectExtBindings :: Set ErasedBinding

instance {-# OVERLAPS #-} Extension '[] where
  collectExtBindings = Set.empty

instance (Extension xs, Material a) => Extension (a : xs) where
  collectExtBindings = Set.unions [Set.singleton $ ErasedBinding (MaterialType $ Proxy @a) (Proxy @a) (Proxy @(B a)), collectExtBindings @(Extends a), collectExtBindings @xs]

collectBindings :: forall m. (Material m) => Set ErasedBinding
collectBindings = Set.union (Set.singleton $ ErasedBinding (MaterialType $ Proxy @m) (Proxy @m) (Proxy @(B m))) $ collectExtBindings @(Extends m)

type ExtensionOf a b = (Material a, IsExtensionOf a b ~ True)

type family IsExtensionOf a b where
  IsExtensionOf a NoMaterial = True
  IsExtensionOf a b = IsExtensionOf' '[a] b

type family IsExtensionOf' (m :: [Type]) (target :: Type) :: Bool where
  IsExtensionOf' '[] a = False
  IsExtensionOf' (a : xs) a = True
  IsExtensionOf' (a : xs) t = Or (IsExtensionOf' (Extends a) t) (IsExtensionOf' xs t)

type family Or a b where
  Or False False = False
  Or a b = True

-- type family ExtendsMaterialRec (e :: Maybe Type) (target :: Type) :: Bool where
--   ExtendsMaterialRec Nothing t = False
--   ExtendsMaterialRec (Just m) t = IsExtensionOf m t

bindings :: forall a. (Material a) => Shader a (B a GPU)
bindings = do
  ShaderState {dict} <- get
  (ErasedValue a) <- pure $ unwrap $ Map.lookup (typeRep $ Proxy @a) dict
  pure $ unsafeCoerce a

extension :: forall e a. (Material e, ExtensionOf a e) => Shader a (B e GPU)
extension = do
  ShaderState {dict} <- get
  (ErasedValue a) <- pure $ unwrap $ Map.lookup (typeRep $ Proxy @e) dict
  pure $ unsafeCoerce a

convertFrag :: forall m. (VO m GPU -> Shader m (ST.Vec 4 ST.Float)) -> VO m GPU -> Shader m (FragmentOutput GPU)
convertFrag s v = do
  s' <- s v
  pure $ FragmentOutput s'

data FetchBindings b where
  FetchBindings ::
    forall x y b.
    (Bindable b) =>
    { camera :: CameraEntity -> Query System x,
      mesh :: MeshEntity -> Query System y,
      result :: x -> y -> b CPU
    } ->
    FetchBindings b

newtype FragmentOutput f = FragmentOutput
  { fragOut1 :: Location f 0 (ST.Vec 4 ST.Float)
  }
  deriving stock (Generic)
  deriving anyclass (ShaderParam)

blendState :: WGPUBlendState
blendState =
  WGPUBlendState
    { color =
        WGPUBlendComponent
          { srcFactor = wGPUBlendFactor_SrcAlpha,
            dstFactor = wGPUBlendFactor_OneMinusSrcAlpha,
            operation = wGPUBlendOperation_Add
          },
      alpha =
        WGPUBlendComponent
          { srcFactor = wGPUBlendFactor_One,
            dstFactor = wGPUBlendFactor_OneMinusSrcAlpha,
            operation = wGPUBlendOperation_Add
          }
    }

data DrawCommand bindings = Draw {binds :: [ErasedBindingValue], vertices :: Nat}

drawCommand :: forall m. (Material m) => RenderDevice -> Ptr WGPURenderPassEncoder -> BindLayout -> DrawCommand (B m) -> IO ()
drawCommand device encoder layout (Draw b vert) = do
  bindGroup <- createBindGroup device layout b
  wgpuRenderPassEncoderSetBindGroup encoder 0 bindGroup 0 (ConstPtr nullPtr)
  wgpuRenderPassEncoderDraw encoder (fromIntegral vert) 1 0 0

render :: forall mat. (Material mat) => RenderDevice -> RenderQueue -> RenderPipeline -> RenderOptions -> Texture -> [DrawCommand (B mat)] -> System ()
render (RenderDevice device) (RenderQueue queue) (RenderPipeline pipeline) RenderOptions {clear} (Texture {texture = output}) commands = liftIO $ do
  bindLayout <- createBindLayout (Set.toList $ collectBindings @mat) (RenderDevice device)

  -- bindGroup <- createBindGroup (RenderDevice device) bindLayout b
  -- RenderPipeline pipeline <- createPipeline @mat (RenderDevice device) (toWGPUFormat format)

  frame <- wgpuTextureCreateView output (ConstPtr nullPtr)

  commandEncoder <- withWGPUString "command encoder" $ \label -> do
    let desc = WGPUCommandEncoderDescriptor {nextInChain = nullPtr, label}
    with desc $ \desc -> wgpuDeviceCreateCommandEncoder device desc

  let (clearValue, loadOp) = case clear of
        Clear (Math.V3 a b c) -> (wgpuColor (realToFrac a) (realToFrac b) (realToFrac c) 1, wGPULoadOp_Clear)
        DontClear -> (wgpuColor 0 0 0 1, wGPULoadOp_Load)

  let attachment =
        WGPURenderPassColorAttachment
          { clearValue,
            storeOp = wGPUStoreOp_Store,
            loadOp,
            resolveTarget = nullPtr,
            depthSlice = wGPU_DEPTH_SLICE_UNDEFINED,
            view = frame,
            nextInChain = nullPtr
          }

  renderPassEncoder <- with attachment $ \attachment -> do
    withWGPUString "render pass encoder" $ \label -> do
      let desc =
            WGPURenderPassDescriptor
              { timestampWrites = ConstPtr nullPtr,
                occlusionQuerySet = nullPtr,
                depthStencilAttachment = ConstPtr nullPtr,
                colorAttachments = ConstPtr attachment,
                colorAttachmentCount = 1,
                label,
                nextInChain = nullPtr
              }

      with desc $ \desc -> do
        wgpuCommandEncoderBeginRenderPass commandEncoder (ConstPtr desc)

  when (renderPassEncoder == nullPtr) $ error "Couldn't encode render pass."

  wgpuRenderPassEncoderSetPipeline renderPassEncoder pipeline
  for_ commands $ drawCommand @mat (RenderDevice device) renderPassEncoder bindLayout
  wgpuRenderPassEncoderEnd renderPassEncoder
  wgpuRenderPassEncoderRelease renderPassEncoder

  commandBuffer <- withWGPUString "command buffer" $ \label -> do
    let desc = WGPUCommandBufferDescriptor {label, nextInChain = nullPtr}
    with desc $ \desc -> wgpuCommandEncoderFinish commandEncoder (ConstPtr desc)

  when (commandBuffer == nullPtr) $ error "Couldn't create command buffer."

  with commandBuffer $ \commandBuffer -> do
    wgpuQueueSubmit queue 1 (ConstPtr commandBuffer)

  wgpuCommandBufferRelease commandBuffer
  wgpuCommandEncoderRelease commandEncoder
  wgpuTextureViewRelease frame

data Clear = DontClear | Clear (Math.V3 Prelude.Float)

data RenderOptions = RenderOptions
  { clear :: Clear,
    nvertices :: Nat
  }

instance Default RenderOptions where
  def =
    RenderOptions
      { clear = DontClear,
        nvertices = 6
      }