{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE UndecidableInstances #-}

-- XXX: TypeError in Compatible generates unused constraint argument
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Engine.Vulkan.Pipeline.Graphics
  ( Config(..)
  , baseConfig

  , Configure
  , Pipeline.Specialization

  , vertexInput
  , formatSize
  , pushPlaceholder

  , Stages(..)
  , stageNames
  , stageFlagBits
  , basicStages
  , vertexOnly
  , StageCode
  , StageSpirv
  , StageReflect

  , Pipeline(..)
  , allocate
  , create
  , Pipeline.destroy

  , bind
  ) where

import RIO

import Data.Bits ((.|.))
import Data.Kind (Type)
import Data.List qualified as List
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import GHC.Generics (Generic1)
import GHC.Stack (callStack, getCallStack, srcLocModule, withFrozenCallStack)
import UnliftIO.Resource (MonadResource, ReleaseKey)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Core12.Promoted_From_VK_EXT_descriptor_indexing qualified as Vk12
import Vulkan.CStruct.Extends (SomeStruct(..), pattern (:&), pattern (::&))
import Vulkan.NamedType ((:::))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (Zero(..))

import Engine.SpirV.Reflect (Reflect)
import Engine.Vulkan.DescSets (Bound(..), Compatible)
import Engine.Vulkan.Pipeline (Pipeline(..))
import Engine.Vulkan.Pipeline qualified as Pipeline
import Engine.Vulkan.Pipeline.Stages (StageInfo(..))
import Engine.Vulkan.Shader qualified as Shader
import Engine.Vulkan.Types (HasVulkan(..), HasRenderPass(..), MonadVulkan, DsBindings, getPipelineCache)
import Render.Code (Code)
import Resource.Collection (Generically1(..))

data Stages a = Stages
  { vert :: a -- ^ vertex
  , tesc :: a -- ^ tessellation control
  , tese :: a -- ^ tessellation evaluation
  , geom :: a -- ^ geometry
  , frag :: a -- ^ fragment
  }
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable, Generic1)
  deriving Applicative via (Generically1 Stages)

instance StageInfo Stages where
  stageNames = Stages
    { vert = "vert"
    , tesc = "tesc"
    , tese = "tese"
    , geom = "geom"
    , frag = "frag"
    }

  stageFlagBits = Stages
    { vert = Vk.SHADER_STAGE_VERTEX_BIT
    , tesc = Vk.SHADER_STAGE_TESSELLATION_CONTROL_BIT
    , tese = Vk.SHADER_STAGE_TESSELLATION_EVALUATION_BIT
    , geom = Vk.SHADER_STAGE_GEOMETRY_BIT
    , frag = Vk.SHADER_STAGE_FRAGMENT_BIT
    }

basicStages
  :: "vert" ::: a
  -> "frag" ::: a
  -> Stages (Maybe a)
basicStages v f = (pure Nothing)
  { vert = Just v
  , frag = Just f
  }

vertexOnly
  :: "vert" ::: a
  -> Stages (Maybe a)
vertexOnly v = (pure Nothing)
  { vert = Just v
  }

type StageCode = Stages (Maybe Code)
type StageSpirv = Stages (Maybe ByteString)
type StageReflect = Reflect Stages

type family Configure pipeline where
  Configure (Pipeline dsl vertices instances) =
    Config
      dsl
      vertices
      instances
      (Pipeline.Specialization (Pipeline dsl vertices instances))

data Config (dsl :: [Type]) vertices instances spec = Config
  { cStages             :: StageSpirv
  , cReflect            :: Maybe StageReflect
  , cVertexInput        :: SomeStruct Vk.PipelineVertexInputStateCreateInfo
  , cDescLayouts        :: Tagged dsl [DsBindings]
  , cPushConstantRanges :: Vector Vk.PushConstantRange
  , cBlend              :: Bool
  , cDepthWrite         :: Bool
  , cDepthTest          :: Bool
  , cDepthCompare       :: Vk.CompareOp
  , cTopology           :: Vk.PrimitiveTopology
  , cCull               :: Vk.CullModeFlagBits
  , cDepthBias          :: Maybe ("constant" ::: Float, "slope" ::: Float)
  , cSpecialization     :: spec
  }

-- | Settings for generic triangle-rendering pipeline.
baseConfig :: Config '[] vertices instances ()
baseConfig = Config
  { cStages             = pure Nothing
  , cVertexInput        = zero
  , cReflect            = Nothing
  , cDescLayouts        = Tagged []
  , cPushConstantRanges = mempty
  , cBlend              = False
  , cDepthWrite         = True
  , cDepthTest          = True
  , cDepthCompare       = Vk.COMPARE_OP_LESS
  , cTopology           = Vk.PRIMITIVE_TOPOLOGY_TRIANGLE_LIST
  , cCull               = Vk.CULL_MODE_BACK_BIT
  , cDepthBias          = Nothing
  , cSpecialization     = ()
  }

-- XXX: consider using instance attrs or uniforms
pushPlaceholder :: Vk.PushConstantRange
pushPlaceholder = Vk.PushConstantRange
  { Vk.stageFlags = Vk.SHADER_STAGE_VERTEX_BIT .|. Vk.SHADER_STAGE_FRAGMENT_BIT
  , Vk.offset     = 0
  , Vk.size       = 4 * dwords
  }
  where
    -- XXX: each 4 word32s eat up one register (on AMD)
    dwords = 4

allocate
  :: ( config ~ Configure pipeline
     , pipeline ~ Pipeline dsl vertices instances
     , spec ~ Pipeline.Specialization pipeline
     , Shader.Specialization spec
     , HasCallStack
     , MonadVulkan env m
     , MonadResource m
     , HasRenderPass renderpass
     )
  => Maybe Vk.Extent2D
  -> Vk.SampleCountFlagBits
  -> Config dsl vertices instances spec
  -> renderpass
  -> m (ReleaseKey, pipeline)
allocate extent msaa config renderpass = withFrozenCallStack do
  ctx <- ask
  Resource.allocate
    (create ctx extent msaa renderpass config)
    (Pipeline.destroy ctx)

create
  :: ( MonadUnliftIO io
     , HasVulkan ctx
     , HasRenderPass renderpass
     , Shader.Specialization spec
     , HasCallStack
     )
  => ctx
  -> Maybe Vk.Extent2D
  -> Vk.SampleCountFlagBits
  -> renderpass
  -> Config dsl vertices instances spec
  -> io (Pipeline dsl vertices instances)
create context mextent msaa renderpass Config{..} = do
  let
    originModule =
      fromString . List.intercalate "|" $
        map (srcLocModule . snd) (getCallStack callStack)

  dsLayouts <- Vector.forM (Vector.fromList $ unTagged cDescLayouts) \bindsFlags -> do
    let
      (binds, flags) = List.unzip bindsFlags

      setCI =
        zero
          { Vk.bindings = Vector.fromList binds
          }
        ::& zero
          { Vk12.bindingFlags = Vector.fromList flags
          }
        :& ()

    Vk.createDescriptorSetLayout device setCI Nothing

  -- TODO: get from outside
  layout <- Vk.createPipelineLayout device (layoutCI dsLayouts) Nothing
  Debug.nameObject device layout originModule

  shader <- Shader.withSpecialization cSpecialization $
    Shader.create context cStages

  let
    cis = Vector.singleton . SomeStruct $
      pipelineCI (Shader.sPipelineStages shader) layout

  Vk.createGraphicsPipelines device cache cis Nothing >>= \case
    (Vk.SUCCESS, pipelines) ->
      case Vector.toList pipelines of
        [one] -> do
          Shader.destroy context shader
          Debug.nameObject device one originModule
          pure Pipeline
            { pipeline     = one
            , pLayout      = Tagged layout
            , pDescLayouts = Tagged dsLayouts
            }
        _ ->
          error "assert: exactly one pipeline requested"
    (err, _) ->
      error $ "createGraphicsPipelines: " <> show err
  where
    device = getDevice context
    cache = getPipelineCache context

    layoutCI dsLayouts = Vk.PipelineLayoutCreateInfo
      { flags              = zero
      , setLayouts         = dsLayouts
      , pushConstantRanges = cPushConstantRanges
      }

    pipelineCI stages layout = zero
      { Vk.stages             = stages
      , Vk.vertexInputState   = Just cVertexInput
      , Vk.inputAssemblyState = Just inputAsembly
      , Vk.viewportState      = Just $ SomeStruct viewportState
      , Vk.rasterizationState = Just $ SomeStruct rasterizationState
      , Vk.multisampleState   = Just $ SomeStruct multisampleState
      , Vk.depthStencilState  = Just depthStencilState
      , Vk.colorBlendState    = Just $ SomeStruct colorBlendState
      , Vk.dynamicState       = dynamicState
      , Vk.layout             = layout
      , Vk.renderPass         = getRenderPass renderpass
      , Vk.subpass            = 0
      , Vk.basePipelineHandle = zero
      }
      where
        inputAsembly = zero
          { Vk.topology               = cTopology
          , Vk.primitiveRestartEnable = restartable
          }

        restartable = elem @Set cTopology
          [ Vk.PRIMITIVE_TOPOLOGY_LINE_STRIP
          , Vk.PRIMITIVE_TOPOLOGY_TRIANGLE_STRIP
          , Vk.PRIMITIVE_TOPOLOGY_TRIANGLE_FAN
          ]

        (viewportState, dynamicState) = case mextent of
          Nothing ->
            ( zero
                { Vk.viewportCount = 1
                , Vk.scissorCount = 1
                }
            , Just zero
                { Vk.dynamicStates = Vector.fromList
                    [ Vk.DYNAMIC_STATE_VIEWPORT
                    , Vk.DYNAMIC_STATE_SCISSOR
                    ]
                }
            )
          Just extent@Vk.Extent2D{width, height} ->
            ( zero
                { Vk.viewports = Vector.fromList
                    [ Vk.Viewport
                        { Vk.x        = 0
                        , Vk.y        = 0
                        , Vk.width    = realToFrac width
                        , Vk.height   = realToFrac height
                        , Vk.minDepth = 0
                        , Vk.maxDepth = 1
                        }
                    ]
                , Vk.scissors = Vector.singleton Vk.Rect2D
                    { Vk.offset = Vk.Offset2D 0 0
                    , extent    = extent
                    }
                }
            , Nothing
            )

        rasterizationState = case cDepthBias of
          Nothing ->
            rasterizationBase
          Just (constantFactor, slopeFactor) ->
            rasterizationBase
              { Vk.depthBiasEnable         = True
              , Vk.depthBiasConstantFactor = constantFactor
              , Vk.depthBiasSlopeFactor    = slopeFactor
              }

        rasterizationBase = zero
          { Vk.depthClampEnable        = False
          , Vk.rasterizerDiscardEnable = False
          , Vk.lineWidth               = 1
          , Vk.polygonMode             = Vk.POLYGON_MODE_FILL
          , Vk.cullMode                = cCull
          , Vk.frontFace               = Vk.FRONT_FACE_CLOCKWISE
          , Vk.depthBiasEnable         = False
          }

        multisampleState = zero
          { Vk.rasterizationSamples = msaa
          , Vk.sampleShadingEnable  = enable
          , Vk.minSampleShading     = if enable then 0.2 else 1.0
          , Vk.sampleMask           = Vector.singleton maxBound
          }
          where
            enable = True -- TODO: check and enable sample rate shading feature

        depthStencilState = zero
          { Vk.depthTestEnable       = cDepthTest
          , Vk.depthWriteEnable      = cDepthWrite
          , Vk.depthCompareOp        = cDepthCompare
          , Vk.depthBoundsTestEnable = False
          , Vk.minDepthBounds        = 0.0 -- Optional
          , Vk.maxDepthBounds        = 1.0 -- Optional
          , Vk.stencilTestEnable     = False
          , Vk.front                 = zero -- Optional
          , Vk.back                  = zero -- Optional
          }

        colorBlendState = zero
          { Vk.logicOpEnable =
              False
          , Vk.attachments = Vector.singleton zero
              { Vk.blendEnable         = cBlend
              , 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_ONE_MINUS_SRC_ALPHA
              , Vk.alphaBlendOp        = Vk.BLEND_OP_ADD
              , Vk.colorWriteMask      = colorRgba
              }
          }

        colorRgba =
          Vk.COLOR_COMPONENT_R_BIT .|.
          Vk.COLOR_COMPONENT_G_BIT .|.
          Vk.COLOR_COMPONENT_B_BIT .|.
          Vk.COLOR_COMPONENT_A_BIT

bind
  :: ( Compatible pipeLayout boundLayout
     , MonadIO m
     )
  => Vk.CommandBuffer
  -> Pipeline pipeLayout vertices instances
  -> Bound boundLayout vertices instances m ()
  -> Bound boundLayout oldVertices oldInstances m ()
bind cb Pipeline{pipeline} (Bound attrAction) = do
  Bound $ Vk.cmdBindPipeline cb Vk.PIPELINE_BIND_POINT_GRAPHICS pipeline
  Bound attrAction

vertexInput :: [(Vk.VertexInputRate, [Vk.Format])] -> SomeStruct Vk.PipelineVertexInputStateCreateInfo
vertexInput bindings = SomeStruct zero
  { Vk.vertexBindingDescriptions   = binds
  , Vk.vertexAttributeDescriptions = attrs
  }
  where
    binds = Vector.fromList do
      (ix, (rate, formats)) <- zip [0..] bindings
      pure Vk.VertexInputBindingDescription
        { binding   = ix
        , stride    = sum $ map formatSize formats
        , inputRate = rate
        }

    attrs = attrBindings $ map snd bindings

-- * Utils

attrBindings :: [[Vk.Format]] -> Vector Vk.VertexInputAttributeDescription
attrBindings bindings = mconcat $ List.unfoldr shiftLocations (0, 0, bindings)
  where
    shiftLocations = \case
      (_binding, _lastLoc, [])           -> Nothing
      (binding, lastLoc, formats : rest) -> Just (bound, next)
        where
          bound = Vector.fromList do
            (ix, format) <- zip [0..] formats
            let offset = sum . map formatSize $ take ix formats
            pure zero
              { Vk.binding  = binding
              , Vk.location = fromIntegral $ lastLoc + ix
              , Vk.format   = format
              , Vk.offset   = offset
              }

          next =
            ( binding + 1
            , lastLoc + Vector.length bound
            , rest
            )

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