{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor #-}
module Types where

import Control.Concurrent.STM.TMVar as STM
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Trans.Resource
import qualified Data.Map.Strict as M
import qualified Data.Vector as V
import Foreign
import Foreign.C.Types (CInt)
import Linear
import qualified SDL
import qualified Vulkan as Vk
import qualified Vulkan.Zero as Vk
import qualified VulkanMemoryAllocator as VMA

data ShaderContainer = ShaderContainer
  { containedVertexShader   :: Maybe Vk.ShaderModule
  , containedFragmentShader :: Maybe Vk.ShaderModule
  }
  deriving (Show)

data EngineData = EngineData
  { engineWindow                 :: SDL.Window
  , engineWindowDimensions       :: V2 CInt
  , enginePhysicalDevice         :: Vk.PhysicalDevice
  , engineLogicalDevice          :: Vk.Device
  , engineInstance               :: Vk.Instance
  , engineSwapchain              :: Vk.SwapchainKHR
  , engineQueue                  :: Vk.Queue
  , engineFramebuffers           :: V.Vector Vk.Framebuffer
  , engineRenderPass             :: Vk.RenderPass
  , engineAllocator              :: VMA.Allocator
  , engineDepthImageView         :: Vk.ImageView
  , engineDepthImage             :: AllocatedImage
  , engineDepthFormat            :: Vk.Format
  , engineFrames                 :: V.Vector FrameData
  , engineGlobalSetLayout        :: Vk.DescriptorSetLayout
  , engineObjectSetLayout        :: Vk.DescriptorSetLayout
  , engineSingleTextureSetLayout :: Vk.DescriptorSetLayout
  , engineDescriptorPool         :: Vk.DescriptorPool
  , engineSceneParameters        :: GPUSceneData
  , engineUploadContext          :: UploadContext
  } 
  deriving (Show)

data AllocatedBuffer = AllocatedBuffer
  { allocatedBuffer  :: Vk.Buffer
  , bufferAllocation :: VMA.Allocation
  , bufferInfo       :: VMA.AllocationInfo
  }
  deriving (Show)

instance Eq AllocatedBuffer where

  (AllocatedBuffer buf1 _ _) == (AllocatedBuffer buf2 _ _) = buf1 == buf2

data Vertex = Vertex
  { vertexPosition :: V3 Float
  , vertexNormal   :: V3 Float
  , vertexColor    :: V4 Float
  , vertexUV       :: V2 Float
  }
  deriving (Show, Eq)

instance Storable Vertex where

  sizeOf _ =
    sizeOf (undefined :: V3 Float) * 2 +
    sizeOf (undefined :: V4 Float) +
    sizeOf (undefined :: V2 Float)

  peek ptr = do
    pos <- peek (castPtr ptr)
    nor <- peek (castPtr ptr `plusPtr` sizeOf pos)
    col <- peek (castPtr ptr `plusPtr` sizeOf pos `plusPtr` sizeOf nor)
    uv <- peek (castPtr ptr `plusPtr` sizeOf pos `plusPtr` sizeOf nor `plusPtr` sizeOf col)
    return $ Vertex pos nor col uv

  poke ptr (Vertex position normal color uv) = do
    let castedV3Ptr = castPtr ptr
    pokeElemOff castedV3Ptr 0 position
    pokeElemOff castedV3Ptr 1 normal
    poke (castPtr (ptr `plusPtr` sizeOf position `plusPtr` sizeOf normal)) color
    poke (castPtr (ptr `plusPtr` sizeOf position `plusPtr` sizeOf normal `plusPtr` sizeOf color)) uv

  alignment _ = 0

class VertexInputDescribable v where

  getVertexDescription :: v -> VertexInputDescription

instance VertexInputDescribable Vertex where

  getVertexDescription v =
    let mainBinding = Vk.zero
          { Vk.binding   = 0
          , Vk.stride    = fromIntegral (sizeOf v)
          , Vk.inputRate = Vk.VERTEX_INPUT_RATE_VERTEX
          } :: Vk.VertexInputBindingDescription
        positionAttribute = Vk.zero
          { Vk.binding  = 0
          , Vk.location = 0
          , Vk.format   = Vk.FORMAT_R32G32B32_SFLOAT
          , Vk.offset   = 0
          } :: Vk.VertexInputAttributeDescription
        normalAttribute = Vk.zero
          { Vk.binding  = 0
          , Vk.location = 1
          , Vk.format   = Vk.FORMAT_R32G32B32_SFLOAT
          , Vk.offset   = fromIntegral (sizeOf (vertexPosition v))
          } :: Vk.VertexInputAttributeDescription
        colorAttribute = Vk.zero
          { Vk.binding  = 0
          , Vk.location = 2
          , Vk.format   = Vk.FORMAT_R32G32B32A32_SFLOAT
          , Vk.offset   = fromIntegral (sizeOf (vertexPosition v) + sizeOf (vertexNormal v))
          } :: Vk.VertexInputAttributeDescription
        uvAttribute = Vk.zero
          { Vk.binding  = 0
          , Vk.location = 3
          , Vk.format   = Vk.FORMAT_R32G32_SFLOAT
          , Vk.offset   = fromIntegral $
            sizeOf (vertexPosition v) + sizeOf (vertexNormal v) + sizeOf (vertexColor v)
          } :: Vk.VertexInputAttributeDescription
    in
      VertexInputDescription
        { vidBindings   = V.fromList [ mainBinding ]
        , vidAttributes = V.fromList
          [ positionAttribute
          , normalAttribute
          , colorAttribute
          , uvAttribute
          ]
        }

data VertexInputDescription = VertexInputDescription
  { vidBindings   :: V.Vector Vk.VertexInputBindingDescription
  , vidAttributes :: V.Vector Vk.VertexInputAttributeDescription
  }
  deriving (Show)

data Mesh = Mesh
  { meshVertices :: V.Vector Vertex
  , meshBuffer   :: AllocatedBuffer
  }
  deriving (Show, Eq)

data MeshPushConstants = MeshPushConstants
  { meshPushData     :: V4 Float
  , meshRenderMatrix :: M44 Float
  }
  deriving (Show)

undefinedMeshPushConstants :: MeshPushConstants
undefinedMeshPushConstants = MeshPushConstants (V4 0 0 0 0) identity

instance Storable MeshPushConstants where

  sizeOf _ = sizeOf (V4 0 0 0 0 :: V4 Float) + sizeOf (identity :: M44 Float)

  peek ptr = do
    dat <- peek (castPtr ptr)
    mat <- peek (castPtr (ptr `plusPtr` sizeOf dat))
    return (MeshPushConstants dat mat)

  poke ptr (MeshPushConstants dat mat) = do
    poke (castPtr ptr) dat
    poke (castPtr $ ptr `plusPtr` sizeOf dat) mat

  alignment _ = 0

data AllocatedImage = AllocatedImage
  { image      :: Vk.Image
  , allocation :: VMA.Allocation
  }
  deriving (Show)

data Material = Material
  { materialTextureSet     :: Vk.DescriptorSet
  , materialPipeline       :: Vk.Pipeline
  , materialPipelineLayout :: Vk.PipelineLayout
  }
  deriving (Show, Eq)

data RenderObject = RenderObject
  { objectMesh     :: String
  , objectMaterial :: String
  , objectMatrix   :: M44 Float
  }
  deriving (Show)

data Texture = Texture
  { textureImage     :: AllocatedImage
  , textureImageView :: Vk.ImageView
  }

data ReadState = ReadState
  { renderables     :: STM.TMVar (V.Vector RenderObject)
  , meshLibrary     :: STM.TMVar (M.Map String Mesh)
  , materialLibrary :: STM.TMVar (M.Map String Material)
  , textureLibrary  :: STM.TMVar (M.Map String Texture)
  }

newtype RenderReader rd m a = RenderReader
  { runRenderInner :: RenderInner rd m a
  }
  deriving
    (Functor, Applicative, Monad, MonadIO, MonadReader rd, MonadResource, MonadFail)

type RenderInner rd m = ReaderT rd m

type Render a = RenderReader ReadState ResIO a

runRender :: ReadState -> Render a -> IO a
runRender rd actions = runResourceT $ flip runReaderT rd $ runRenderInner actions

data FrameData = FrameData
  { framePresentSemaphore  :: Vk.Semaphore
  , frameRenderSemaphore   :: Vk.Semaphore
  , frameRenderFence       :: Vk.Fence
  , frameCommandPool       :: Vk.CommandPool
  , frameMainCommandBuffer :: Vk.CommandBuffer
  , frameCameraBuffer      :: AllocatedBuffer
  , frameSceneBuffer       :: AllocatedBuffer
  , frameObjectBuffer      :: AllocatedBuffer
  , frameIndirectBuffer    :: AllocatedBuffer
  , frameGlobalDescriptor  :: Vk.DescriptorSet
  , frameObjectDescriptor  :: Vk.DescriptorSet
  }
  deriving (Show)

newtype GPUObjectData = GPUObjectData
  { objectModelMatrix :: M44 Float
  }

undefinedGPUObjectData :: GPUObjectData
undefinedGPUObjectData = GPUObjectData identity

instance Storable GPUObjectData where

  sizeOf (GPUObjectData _) = sizeOf (identity :: M44 Float)

  alignment _ = 0

  peek ptr = do
    GPUObjectData <$> peek (castPtr ptr)

  poke ptr (GPUObjectData modelMatrix) =
    poke (castPtr ptr) $ transpose modelMatrix

data GPUCameraData = GPUCameraData
  { view           :: M44 Float
  , projection     :: M44 Float
  , viewProjection :: M44 Float
  }

instance Storable GPUCameraData where

  sizeOf _ =
    3 * sizeOf (undefined :: M44 Float)

  alignment _ = 0

  peek ptr = do
    v <- peek (castPtr ptr)
    p <- peek (castPtr ptr `plusPtr` sizeOf v)
    vp <- peek (castPtr ptr `plusPtr` sizeOf v `plusPtr` sizeOf p)
    return $ GPUCameraData v p vp

  poke ptr (GPUCameraData pview pprojection pviewProjection) = do
    poke (castPtr ptr) $ transpose pview
    poke (castPtr ptr `plusPtr` sizeOf pview) $ transpose pprojection
    poke (castPtr ptr `plusPtr` sizeOf pview `plusPtr` sizeOf pprojection) $
      transpose pviewProjection

data GPUSceneData = GPUSceneData
  { fogColor     :: V4 Float
  , fogDistance  :: V4 Float
  , ambientColor :: V4 Float
  , sunDirection :: V4 Float
  , sunColor     :: V4 Float
  }
  deriving (Show)

undefinedGPUSceneData :: GPUSceneData
undefinedGPUSceneData = GPUSceneData
  (V4 0 0 0 0)
  (V4 0 0 0 0)
  (V4 0 0 0 0)
  (V4 0 0 0 0)
  (V4 0 0 0 0)

instance Storable GPUSceneData where

  sizeOf _ =
    5 * sizeOf (V4 0 0 0 0 :: V4 Float)

  alignment _ = 0

  peek ptr = do
    fogCol <- peek (castPtr ptr)
    fogDst <- peek (castPtr $ ptr `plusPtr` sizeOf fogCol)
    ambCol <- peek (castPtr $ ptr `plusPtr` (sizeOf fogCol + sizeOf fogDst))
    sunDir <- peek (castPtr $ ptr `plusPtr` (sizeOf fogCol + sizeOf fogDst + sizeOf ambCol))
    sunCol <- peek (castPtr $ ptr `plusPtr` (sizeOf fogCol + sizeOf fogDst + sizeOf ambCol + sizeOf sunDir))
    return $ GPUSceneData fogCol fogDst ambCol sunDir sunCol

  poke ptr (GPUSceneData fogCol fogDist ambCol sunDir sunCol) = do
    poke (castPtr ptr) fogCol
    poke (castPtr $ ptr `plusPtr` sizeOf fogCol) fogDist
    poke (castPtr $ ptr `plusPtr` (sizeOf fogCol + sizeOf fogDist)) ambCol
    poke (castPtr $ ptr `plusPtr` (sizeOf fogCol + sizeOf fogDist + sizeOf ambCol)) sunDir
    poke (castPtr $ ptr `plusPtr` (sizeOf fogCol + sizeOf fogDist + sizeOf ambCol + sizeOf sunDir)) sunCol

data UploadContext = UploadContext
  { uploadFence         :: Vk.Fence
  , uploadCommandPool   :: Vk.CommandPool
  , uploadCommandBuffer :: Vk.CommandBuffer
  }
  deriving (Show)

data IndirectBatch = IndirectBatch
  { batchMesh     :: Mesh
  , batchMaterial :: Material
  , batchFirst    :: Word32
  , batchCount    :: Word32
  }
