{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module GraphicsPipeline where

import Linear
import qualified Control.Concurrent.STM as STM
import Control.Monad.Reader
import Control.Monad.Trans.Resource
import Data.Bits
-- import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS
import qualified Data.Map.Strict as M
import qualified Data.Vector as V
import Foreign (sizeOf)
import Foreign.C.Types (CInt)
import qualified Vulkan as VK
import qualified Vulkan.Core10 as Vk
import qualified Vulkan.Zero as Vk
import qualified Vulkan.CStruct.Extends as Vk
import qualified Vulkan.Utils.ShaderQQ.GLSL.Shaderc as Vk

-- internal imports
import Types

loadShader
  :: (MonadResource m)
  => Vk.Device
  -> FilePath
  -> String
  -> m Vk.ShaderModule
loadShader logicalDevice shaderPath stage = do
  shaderSource <- liftIO $ BS.readFile shaderPath

  (warnings, eShader) <- Vk.compileShader Nothing Nothing stage (BS.unpack shaderSource)
  unless (null warnings) $ liftIO $ do
    putStrLn "loadShader: Shader compilation warnings:"
    print warnings
  shader <- either
    (\errors -> do
      error ("loadShader: Compilation errors:" <> "\n" <> show errors)
      )
    return
    eShader


  let createInfo = Vk.zero
        { Vk.code = shader
        }

  snd <$>
    Vk.withShaderModule logicalDevice createInfo Nothing allocate

createRenderPass
  :: (MonadResource m)
  => Vk.Device
  -> Vk.Format
  -> Vk.Format
  -> m Vk.RenderPass
createRenderPass
  logicalDevice
  swapchainImageFormat
  depthFormat
  = do

    let colorAttachmentDescription = (Vk.zero :: Vk.AttachmentDescription)
          { Vk.format         = swapchainImageFormat
          , Vk.samples        = Vk.SAMPLE_COUNT_1_BIT
          , Vk.loadOp         = Vk.ATTACHMENT_LOAD_OP_CLEAR
          , Vk.storeOp        = Vk.ATTACHMENT_STORE_OP_STORE
          , Vk.stencilLoadOp  = Vk.ATTACHMENT_LOAD_OP_DONT_CARE
          , Vk.stencilStoreOp = Vk.ATTACHMENT_STORE_OP_DONT_CARE
          , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
          , Vk.finalLayout    = Vk.IMAGE_LAYOUT_PRESENT_SRC_KHR
          }
        colorAttachmentReference = (Vk.zero :: Vk.AttachmentReference)
          { Vk.attachment = 0
          , Vk.layout     = Vk.IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
          }
        depthAttachmentDescription = (Vk.zero :: Vk.AttachmentDescription)
          { Vk.flags          = bit 0
          , Vk.format         = depthFormat
          , Vk.samples        = Vk.SAMPLE_COUNT_1_BIT
          , Vk.loadOp         = Vk.ATTACHMENT_LOAD_OP_CLEAR
          , Vk.storeOp        = Vk.ATTACHMENT_STORE_OP_STORE
          , Vk.stencilLoadOp  = Vk.ATTACHMENT_LOAD_OP_CLEAR
          , Vk.stencilStoreOp = Vk.ATTACHMENT_STORE_OP_DONT_CARE
          , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
          , Vk.finalLayout    = Vk.IMAGE_LAYOUT_DEPTH_READ_ONLY_STENCIL_ATTACHMENT_OPTIMAL
          }
        depthAttachmentReference = (Vk.zero :: Vk.AttachmentReference)
          { Vk.attachment = 1
          , Vk.layout     = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
          }
        subpassDescriptor = (Vk.zero :: Vk.SubpassDescription)
          { Vk.pipelineBindPoint      = Vk.PIPELINE_BIND_POINT_GRAPHICS
          , Vk.colorAttachments       = V.singleton colorAttachmentReference
          , Vk.depthStencilAttachment = Just depthAttachmentReference
          }
        subpassDependency = Vk.zero
          { Vk.srcSubpass    = Vk.SUBPASS_EXTERNAL
          , Vk.dstSubpass    = 0
          , Vk.srcStageMask  = Vk.PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT
          , Vk.dstStageMask  = Vk.PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT
          , Vk.srcAccessMask = Vk.ACCESS_MEMORY_READ_BIT
          , Vk.dstAccessMask = Vk.ACCESS_COLOR_ATTACHMENT_WRITE_BIT
          }
        subpassDepthDependency = Vk.zero
          { Vk.srcSubpass    = Vk.SUBPASS_EXTERNAL
          , Vk.dstSubpass    = 0
          , Vk.srcStageMask  =
            Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT .|. Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT
          , Vk.srcAccessMask = Vk.ACCESS_MEMORY_READ_BIT
          , Vk.dstStageMask  =
            Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT .|. Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT
          , Vk.dstAccessMask = Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
          }
        renderPassCreateInfo = Vk.zero
          { Vk.attachments  = V.fromList [ colorAttachmentDescription, depthAttachmentDescription ]
          , Vk.subpasses    = V.singleton subpassDescriptor
          , Vk.dependencies = V.fromList [ subpassDependency, subpassDepthDependency ]
          } :: Vk.RenderPassCreateInfo '[]

    snd <$>
      Vk.withRenderPass logicalDevice renderPassCreateInfo Nothing allocate

createShaderStageCreateInfo
  :: Vk.ShaderStageFlagBits
  -> Vk.ShaderModule
  -> Vk.PipelineShaderStageCreateInfo '[]
createShaderStageCreateInfo stageBit shaderModule =
  Vk.zero
    { Vk.stage   = stageBit
    , Vk.module' = shaderModule
    , Vk.name    = "main"
    }

createVertexInputAssemblyStateCreateInfo
  :: Vk.PrimitiveTopology
  -> Vk.PipelineInputAssemblyStateCreateInfo
createVertexInputAssemblyStateCreateInfo topology =
  Vk.zero
    { Vk.topology               = topology
    , Vk.primitiveRestartEnable = False
    }

createRasterizationStateCreateInfo
  :: Vk.PolygonMode
  -> Vk.SomeStruct Vk.PipelineRasterizationStateCreateInfo
createRasterizationStateCreateInfo mode =
  Vk.SomeStruct $ Vk.zero
    { Vk.depthClampEnable        = False
    , Vk.rasterizerDiscardEnable = False
    , Vk.polygonMode             = mode
    , Vk.lineWidth               = 1
    , Vk.cullMode                = Vk.CULL_MODE_NONE
    , Vk.frontFace               = Vk.FRONT_FACE_CLOCKWISE
    , Vk.depthBiasEnable         = False
    , Vk.depthBiasConstantFactor = 0
    , Vk.depthBiasClamp          = 0
    , Vk.depthBiasSlopeFactor    = 0
    }

createDepthStencilStateCreateInfo
  :: Bool
  -> Bool
  -> Vk.CompareOp
  -> Vk.PipelineDepthStencilStateCreateInfo
createDepthStencilStateCreateInfo depthTest depthWrite compareOperation =
  Vk.zero
    { Vk.depthTestEnable       = depthTest
    , Vk.depthWriteEnable      = depthWrite
    , Vk.depthCompareOp        = compareOperation
    , Vk.depthBoundsTestEnable = False
    , Vk.minDepthBounds        = 0
    , Vk.maxDepthBounds        = 1
    , Vk.stencilTestEnable     = False
    }

createMultisampleStateCreateInfo
  :: Vk.SomeStruct Vk.PipelineMultisampleStateCreateInfo
createMultisampleStateCreateInfo =
  Vk.SomeStruct $ Vk.zero
    { Vk.sampleShadingEnable  = False
    , Vk.rasterizationSamples = Vk.SAMPLE_COUNT_1_BIT
    }

createGraphicsPipelines
  :: (MonadResource m, MonadReader ReadState m)
  => Vk.Device
  -> Vk.RenderPass
  -> ShaderContainer
  -> V2 CInt
  -> Vk.PipelineLayout
  -> Maybe Vk.PipelineDepthStencilStateCreateInfo
  -> String
  -> m ()
createGraphicsPipelines
  logicalDevice
  renderPass
  shaderContainer
  (V2 width height)
  pipelineLayout
  depthState
  materialName
  = do
    meshLib <- (liftIO . STM.atomically . STM.readTMVar) =<< asks meshLibrary
    let pipelineStagesCreateInfos =
          V.fromList $ map Vk.SomeStruct
            ( maybe
              []
              (\shader ->
                [ createShaderStageCreateInfo
                  Vk.SHADER_STAGE_VERTEX_BIT
                  shader
                ]
              )
              (containedVertexShader shaderContainer)
              ++
              maybe
              []
              (\shader ->
                [ createShaderStageCreateInfo
                  Vk.SHADER_STAGE_FRAGMENT_BIT
                  shader
                ]
              )
              (containedFragmentShader shaderContainer)
            )
        dynamicStates =
          V.fromList
            [ Vk.DYNAMIC_STATE_VIEWPORT
            , Vk.DYNAMIC_STATE_SCISSOR
            ]
        pipelineDynamicStateCreateInfo = Vk.zero
          { Vk.dynamicStates = dynamicStates
          }
        vertexDescription = getVertexDescription (V.head $ meshVertices $ meshLib M.! "mask")
        pipelineVertexInputCreateInfo = Vk.zero
          { Vk.vertexBindingDescriptions   = vidBindings vertexDescription
          , Vk.vertexAttributeDescriptions = vidAttributes vertexDescription
          }
        viewport = Vk.Viewport
          { Vk.x        = 0
          , Vk.y        = 0
          , Vk.width    = fromIntegral width
          , Vk.height   = fromIntegral height
          , Vk.minDepth = 0
          , Vk.maxDepth = 1
          }
        scissor = Vk.Rect2D
          { Vk.offset = Vk.Offset2D 0 0
          , Vk.extent = Vk.Extent2D (fromIntegral width) (fromIntegral height)
          }
        pipelineViewportStateCreateInfo = Vk.zero
          { Vk.viewports = V.singleton viewport
          , Vk.scissors  = V.singleton scissor
          }
        colorBlendAttachment = Vk.zero
          { Vk.colorWriteMask      =
            Vk.COLOR_COMPONENT_R_BIT .|.
            Vk.COLOR_COMPONENT_G_BIT .|.
            Vk.COLOR_COMPONENT_B_BIT .|.
            Vk.COLOR_COMPONENT_A_BIT
          , Vk.blendEnable         = True
          , Vk.srcColorBlendFactor = Vk.BLEND_FACTOR_ONE
          , Vk.dstColorBlendFactor = Vk.BLEND_FACTOR_ONE_MINUS_SRC_ALPHA
          , Vk.colorBlendOp        = Vk.BLEND_OP_ADD
          , Vk.srcAlphaBlendFactor = Vk.BLEND_FACTOR_ONE
          , Vk.dstAlphaBlendFactor = Vk.BLEND_FACTOR_ZERO
          , Vk.alphaBlendOp        = Vk.BLEND_OP_ADD
          }
        pipelineColorBlendStateCreateInfo = Vk.zero
          { Vk.logicOpEnable = False
          , Vk.attachments = V.singleton colorBlendAttachment
          }
        pipelineCreateInfo = Vk.zero
          { Vk.stageCount         = fromIntegral (V.length pipelineStagesCreateInfos)
          , Vk.stages             = pipelineStagesCreateInfos
          , Vk.vertexInputState   = Just $ Vk.SomeStruct pipelineVertexInputCreateInfo
          , Vk.inputAssemblyState =
            Just $ createVertexInputAssemblyStateCreateInfo VK.PRIMITIVE_TOPOLOGY_TRIANGLE_LIST
          , Vk.viewportState      = Just $ Vk.SomeStruct pipelineViewportStateCreateInfo
          , Vk.rasterizationState = Just $ createRasterizationStateCreateInfo Vk.POLYGON_MODE_FILL
          , Vk.multisampleState   = Just createMultisampleStateCreateInfo
          , Vk.colorBlendState    = Just $ Vk.SomeStruct pipelineColorBlendStateCreateInfo
          , Vk.dynamicState       = Just pipelineDynamicStateCreateInfo
          , Vk.depthStencilState  = depthState
          , Vk.layout             = pipelineLayout
          , Vk.renderPass         = renderPass
          , Vk.subpass            = 0
          , Vk.basePipelineHandle = Vk.NULL_HANDLE
          , Vk.basePipelineIndex  = -1
          }

    pipeline <-
      (do
        (result, pipelines) <- snd <$> Vk.withGraphicsPipelines
          logicalDevice
          Vk.NULL_HANDLE
          (V.singleton (Vk.SomeStruct pipelineCreateInfo))
          Nothing
          allocate
        unless (result == Vk.SUCCESS) $
          error "createGraphicsPipelines: Failed creating pipelines"
        return pipelines
        )
    let material = Material Vk.NULL_HANDLE (V.head pipeline) pipelineLayout
    matLibraryTMVar <- asks materialLibrary
    matLibrary <- liftIO $ STM.atomically $ STM.readTMVar matLibraryTMVar
    let newMatLibrary = M.insert materialName material matLibrary
    void $ liftIO $ STM.atomically $ STM.swapTMVar matLibraryTMVar newMatLibrary

createMeshPipelineLayout
  :: MonadResource m
  => Vk.Device
  -> V.Vector Vk.DescriptorSetLayout
  -> m Vk.PipelineLayout
createMeshPipelineLayout logicalDevice layouts = do
  let pushConstantRange = Vk.zero
        { Vk.offset     = 0
        , Vk.size       = fromIntegral (sizeOf undefinedMeshPushConstants)
        , Vk.stageFlags = Vk.SHADER_STAGE_VERTEX_BIT
        }
      pipelineLayoutCreateInfo = Vk.zero
        { Vk.pushConstantRanges = V.singleton pushConstantRange
        , Vk.setLayouts         = layouts
        }

  snd <$>
    Vk.withPipelineLayout logicalDevice pipelineLayoutCreateInfo Nothing allocate
