Core.hs

module Mischief.Render.Core where

import Control.Monad
import Control.Monad.IO.Class
import Data.ByteString qualified as BS
import Data.Default (Default)
import Data.Maybe (isNothing)
import Data.Primitive (Ptr)
import Data.Primitive.Ptr (nullPtr)
import Foreign (Storable (peek, poke), alloca, castPtr, free, malloc, with)
import Foreign.C (peekCString)
import Foreign.C.ConstPtr
import Mischief.ECS
import Mischief.ECS.Events
import Mischief.ECS.Hooks
import Mischief.ECS.Systems qualified as Systems
import Mischief.Render.TextureFormat
import Mischief.SDL (SDLPlugin (..), SDLWindow (SDLWindow))
import Mischief.SDL.Window
import Mischief.WGPU
import Mischief.WGPU.Opaque
import Mischief.WGPU.Types.Enums
import Mischief.WGPU.Types.General
import SDL3.Sys qualified as SDL3
import SDL3.Sys.Bindgen.Metal.FunPtr qualified as SDL3

type CameraEntity = Entity

type MeshEntity = Entity

newtype RenderInstance = RenderInstance (Ptr WGPUInstance) deriving anyclass (Component)

newtype RenderSurface = RenderSurface (Ptr WGPUSurface)

newtype DefaultSampler = DefaultSampler Sampler deriving anyclass (Component)

instance Component RenderSurface where
  onRemove =
    [ hook $ \(HookContext entity) -> do
        Just (RenderSurface surface) <- single [q|entity. RenderSurface|]
        liftIO $ wgpuSurfaceRelease surface
    ]

newtype RenderDevice = RenderDevice (Ptr WGPUDevice)

instance Component RenderDevice where
  onRemove =
    [ hook $ \(HookContext entity) -> do
        Just (RenderDevice device) <- single [q|entity. RenderDevice|]
        liftIO $ wgpuDeviceRelease device
    ]

newtype RenderAdapter = RenderAdapter (Ptr WGPUAdapter)

instance Component RenderAdapter where
  onRemove =
    [ hook $ \(HookContext entity) -> do
        Just (RenderAdapter adapter) <- single [q|entity. RenderAdapter|]
        liftIO $ wgpuAdapterRelease adapter
    ]

newtype RenderQueue = RenderQueue (Ptr WGPUQueue)

instance Component RenderQueue where
  onRemove =
    [ hook $ \(HookContext entity) -> do
        Just (RenderQueue queue) <- single [q|entity. RenderQueue|]
        liftIO $ wgpuQueueRelease queue
    ]

newtype WindowFormat = WindowFormat WGPUTextureFormat deriving anyclass (Component)

newtype WindowAlphaMode = WindowAlphaMode WGPUCompositeAlphaMode deriving anyclass (Component)

-- | Creates a WGPU Surface, Adapter, and Device for each new window. Inserts them as components.
onAddWindow :: OnAdd SDLWindow -> System ()
onAddWindow (OnAdd entity) = do
  Just (SDLWindow window, WindowSize w h) <- single [q|entity. SDLWindow, WindowSize|]

  windowProps <- liftIO $ SDL3.getWindowProperties window
  driver <- liftIO (peekCString . unConstPtr =<< SDL3.getCurrentVideoDriver)

  Just (RenderInstance wgpuInstance) <- res @RenderInstance

  surface <- liftIO $ getSurfaceDescriptor wgpuInstance driver windowProps window

  when (surface == nullPtr) $ error "Couldn't create WGPU surface."

  tryAdapter <- res @RenderAdapter
  (adapter, device, _) <- case tryAdapter of
    Just (RenderAdapter adapter) -> do
      Just (RenderDevice device) <- res @RenderDevice
      Just (RenderQueue queue) <- res @RenderQueue
      pure (adapter, device, queue)
    Nothing -> do
      adapter <- liftIO $ wgpuInstanceRequestAdapter wgpuInstance surface
      device <- liftIO $ wgpuAdapterRequestDevice adapter
      queue <- liftIO $ wgpuDeviceGetQueue device

      insertRes (RenderAdapter adapter)
      insertRes (RenderDevice device)
      insertRes (RenderQueue queue)
      pure (adapter, device, queue)

  surfaceCapabilities <- liftIO $ malloc @WGPUSurfaceCapabilities
  liftIO $ wgpuSurfaceGetCapabilities surface adapter surfaceCapabilities

  cap <- liftIO $ peek surfaceCapabilities
  liftIO $ free surfaceCapabilities
  let (ConstPtr formats) = cap.formats
  let (ConstPtr alphaModes) = cap.alphaModes
  format <- liftIO $ peek formats
  alphaMode <- liftIO $ peek alphaModes

  let config =
        newWGPUSurfaceConfiguration
          { device,
            usage = wGPUTextureUsage_RenderAttachment,
            format,
            presentMode = wGPUPresentMode_Fifo,
            alphaMode,
            width = fromIntegral w,
            height = fromIntegral h
          }

  liftIO $ with config $ \config -> do
    wgpuSurfaceConfigure surface (ConstPtr config)

  insert (RenderSurface surface, WindowFormat format, WindowAlphaMode alphaMode) entity

  sampler <- newSampler (RenderDevice device)
  insertRes $ DefaultSampler sampler

onSetWindowSize :: OnSet WindowSize -> System ()
onSetWindowSize (OnSet window) = query_ $ do
  (WindowSize w h, RenderSurface surface, WindowAlphaMode alphaMode, WindowFormat format) <- [q|window. WindowSize, RenderSurface, WindowAlphaMode, WindowFormat|]
  (RenderDevice device) <- qres

  let config =
        newWGPUSurfaceConfiguration
          { device,
            usage = wGPUTextureUsage_RenderAttachment,
            format,
            presentMode = wGPUPresentMode_Fifo,
            alphaMode,
            width = fromIntegral w,
            height = fromIntegral h
          }

  qlift $ do
    liftIO $ with config $ \config -> do
      wgpuSurfaceConfigure surface (ConstPtr config)

-- | Maps a SDL driver name and window to the correct surface texture for WGPU.
--
-- TODO: implement the rest of the drivers.
getSurfaceDescriptor :: Ptr WGPUInstance -> String -> SDL3.SDL_PropertiesID -> Ptr SDL3.SDL_Window -> IO (Ptr WGPUSurface)
getSurfaceDescriptor inst "x11" props _ = do
  display <- BS.useAsCString SDL3.sDL_PROP_WINDOW_X11_DISPLAY_POINTER $ \b -> SDL3.getPointerProperty props (ConstPtr b) nullPtr
  window <- BS.useAsCString SDL3.sDL_PROP_WINDOW_X11_WINDOW_NUMBER $ \b -> SDL3.getNumberProperty props (ConstPtr b) (-1)

  when (display == nullPtr || window == -1) $ error "Can't obtain x11 window."

  let chain = WGPUChainedStruct {next = nullPtr, sType = wGPUSType_SurfaceSourceXlibWindow}
  let xlib = WGPUSurfaceSourceXlibWindow {chain, display, window}

  with xlib $ \xlib ->
    with
      ( WGPUSurfaceDescriptor
          { nextInChain = castPtr xlib,
            label = WGPUStringView {_data = ConstPtr nullPtr, length = 0}
          }
      )
      $ wgpuInstanceCreateSurface inst . ConstPtr
getSurfaceDescriptor inst "windows" props _ = do
  hwnd <- BS.useAsCString SDL3.sDL_PROP_WINDOW_WIN32_HWND_POINTER $ \b -> SDL3.getPointerProperty props (ConstPtr b) nullPtr
  hinstance <- BS.useAsCString SDL3.sDL_PROP_WINDOW_WIN32_INSTANCE_POINTER $ \b -> SDL3.getPointerProperty props (ConstPtr b) nullPtr

  let chain = WGPUChainedStruct {next = nullPtr, sType = wGPUSType_SurfaceSourceWindowsHWND}
  let win = WGPUSurfaceSourceWindowsHWND {chain, hinstance, hwnd}

  with win $ \win ->
    with
      ( WGPUSurfaceDescriptor
          { nextInChain = castPtr win,
            label = WGPUStringView {_data = ConstPtr nullPtr, length = 0}
          }
      )
      $ wgpuInstanceCreateSurface inst . ConstPtr
getSurfaceDescriptor inst "cocoa" _ window = do
  x <- SDL3.metalCreateView window
  layer <- SDL3.metalGetLayer x

  when (layer == nullPtr) $ error "Couldn't obtain metal layer!"

  let chain = WGPUChainedStruct {next = nullPtr, sType = wGPUSType_SurfaceSourceMetalLayer}
  let cocoa = WGPUSurfaceSourceMetalLayer {chain, layer}

  with cocoa $ \cocoa ->
    with
      ( WGPUSurfaceDescriptor
          { nextInChain = castPtr cocoa,
            label = WGPUStringView {_data = ConstPtr nullPtr, length = 0}
          }
      )
      $ wgpuInstanceCreateSurface inst . ConstPtr
getSurfaceDescriptor _ s _ _ = error $ "Incompatible driver: " ++ s

newtype BindLayout = BindLayout (Ptr WGPUBindGroupLayout)

newtype RenderPipeline = RenderPipeline (Ptr WGPURenderPipeline)

-- newtype TextureFormat = TextureFormat WGPUTextureFormat

newtype Sampler = Sampler (Ptr WGPUSampler)

getRenderingResources :: System (Maybe (RenderAdapter, RenderDevice, RenderQueue))
getRenderingResources = do
  adapter <- res @RenderAdapter
  device <- res @RenderDevice
  queue <- res @RenderQueue
  pure $ (,,) <$> adapter <*> device <*> queue

data RenderFirst = RenderFirst deriving (Schedule)

data RenderPrepare = RenderPrepare deriving (Schedule)

data RenderUpdate = RenderUpdate deriving (Schedule)

data RenderTransparent = RenderTransparent deriving (Schedule)

data RenderPostprocess = RenderPostprocess deriving (Schedule)

data RenderLast = RenderLast deriving (Schedule)

getFormat :: RenderSurface -> RenderAdapter -> System TextureFormat
getFormat (RenderSurface surface) (RenderAdapter adapter) = liftIO $ do
  surfaceCapabilities <- malloc @WGPUSurfaceCapabilities
  wgpuSurfaceGetCapabilities surface adapter surfaceCapabilities

  cap <- peek surfaceCapabilities
  free surfaceCapabilities
  let (ConstPtr formats) = cap.formats
  format <- peek formats

  pure $ fromWGPUFormat format

newSampler :: RenderDevice -> System Sampler
newSampler (RenderDevice device) = liftIO $ do
  let samplerDesc =
        WGPUSamplerDescriptor
          { addressModeU = wGPUAddressMode_ClampToEdge,
            addressModeV = wGPUAddressMode_ClampToEdge,
            addressModeW = wGPUAddressMode_ClampToEdge,
            magFilter = wGPUFilterMode_Linear,
            minFilter = wGPUFilterMode_Linear,
            mipmapFilter = wGPUMipmapFilterMode_Nearest,
            nextInChain = nullPtr,
            label = WGPUStringView (ConstPtr nullPtr) 0,
            lodMinClamp = 0,
            lodMaxClamp = 32,
            compare = wGPUCompareFunction_Undefined,
            maxAnisotropy = 1
          }
  Sampler <$> with samplerDesc (wgpuDeviceCreateSampler device . ConstPtr)