{-# LANGUAGE AllowAmbiguousTypes #-}

module Engine.Vulkan.Format
  ( HasVkFormat(..)
  , genericVkFormat

  , formatSize
  ) where

import RIO
import GHC.Generics
import Geomancy

import Data.Kind (Constraint, Type)
import Geomancy.IVec3 qualified as IVec3
import Geomancy.UVec3 qualified as UVec3
import Geomancy.Vec3 qualified as Vec3
import Vulkan.Core10 qualified as Vk

class HasVkFormat a where
  getVkFormat :: [Vk.Format]

  default getVkFormat
    :: GVkFormat (Rep a)
    => [Vk.Format]
  getVkFormat = genericVkFormat @a

instance HasVkFormat () where
  getVkFormat = []

instance HasVkFormat Float where
  getVkFormat = [Vk.FORMAT_R32_SFLOAT]

instance HasVkFormat Vec2 where
  getVkFormat = [Vk.FORMAT_R32G32_SFLOAT]

instance HasVkFormat Vec3 where
  getVkFormat = [Vk.FORMAT_R32G32B32A32_SFLOAT]

instance HasVkFormat Vec3.Packed where
  getVkFormat = [Vk.FORMAT_R32G32B32_SFLOAT]

instance HasVkFormat Vec4 where
  getVkFormat = [Vk.FORMAT_R32G32B32A32_SFLOAT]

instance HasVkFormat Quaternion where
  getVkFormat = [Vk.FORMAT_R32G32B32A32_SFLOAT]

instance HasVkFormat Mat4 where
  getVkFormat =
    [ Vk.FORMAT_R32G32B32A32_SFLOAT
    , Vk.FORMAT_R32G32B32A32_SFLOAT
    , Vk.FORMAT_R32G32B32A32_SFLOAT
    , Vk.FORMAT_R32G32B32A32_SFLOAT
    ]

instance HasVkFormat Transform where
  getVkFormat = getVkFormat @Mat4

instance HasVkFormat Int32 where
  getVkFormat = [Vk.FORMAT_R32_SINT]

instance HasVkFormat IVec2 where
  getVkFormat = [Vk.FORMAT_R32G32_SINT]

instance HasVkFormat IVec3 where
  getVkFormat = [Vk.FORMAT_R32G32B32_SINT]

instance HasVkFormat IVec3.Packed where
  getVkFormat = [Vk.FORMAT_R32G32B32_SINT]

instance HasVkFormat IVec4 where
  getVkFormat = [Vk.FORMAT_R32G32B32A32_SINT]

instance HasVkFormat Word32 where
  getVkFormat = [Vk.FORMAT_R32_UINT]

instance HasVkFormat UVec2 where
  getVkFormat = [Vk.FORMAT_R32G32_UINT]

instance HasVkFormat UVec3 where
  getVkFormat = [Vk.FORMAT_R32G32B32A32_UINT]

instance HasVkFormat UVec3.Packed where
  getVkFormat = [Vk.FORMAT_R32G32B32_UINT]

instance HasVkFormat UVec4 where
  getVkFormat = [Vk.FORMAT_R32G32B32A32_UINT]

instance HasVkFormat v => HasVkFormat (Point v) where
  getVkFormat = getVkFormat @v

genericVkFormat
  :: forall a
  .  GVkFormat (Rep a)
  => [Vk.Format]
genericVkFormat = gVkFormat (Proxy @(Rep a))

type GVkFormat :: (Type -> Type) -> Constraint
class GVkFormat f where
  gVkFormat :: proxy f -> [Vk.Format]

instance GVkFormat f => GVkFormat (M1 c cb f) where
  gVkFormat _m1 = gVkFormat (Proxy @f)

instance (GVkFormat l, GVkFormat r) => GVkFormat (l :*: r) where
  gVkFormat _lr = gVkFormat (Proxy @l) <> gVkFormat (Proxy @r)

instance HasVkFormat a => GVkFormat (K1 r a) where
  gVkFormat _k1 = getVkFormat @a

formatSize :: Integral a => Vk.Format -> a
formatSize = \case
  Vk.FORMAT_R32G32B32A32_SFLOAT -> 16
  Vk.FORMAT_R32G32B32_SFLOAT    -> 12
  Vk.FORMAT_R32G32_SFLOAT       -> 8
  Vk.FORMAT_R32_SFLOAT          -> 4

  Vk.FORMAT_R32G32B32A32_UINT -> 16
  Vk.FORMAT_R32G32B32_UINT    -> 12
  Vk.FORMAT_R32G32_UINT       -> 8
  Vk.FORMAT_R32_UINT          -> 4

  Vk.FORMAT_R32G32B32A32_SINT -> 16
  Vk.FORMAT_R32G32B32_SINT    -> 12
  Vk.FORMAT_R32G32_SINT       -> 8
  Vk.FORMAT_R32_SINT          -> 4

  format ->
    error $ "Format size unknown: " <> show format