{-# LANGUAGE AllowAmbiguousTypes #-}
{-# 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

  , bind

  , HasVertexInputBindings(..)
  , vertexFormat
  , instanceFormat
  ) where

import RIO
import GHC.Generics

import Data.Bits ((.|.))
import Data.Kind (Type)
import Data.List qualified as List
import Data.Tagged (Tagged(..))
import Data.Type.Equality (type (~))
import Data.Vector qualified as Vector
import Geomancy (Transform)
import GHC.Stack (withFrozenCallStack)
import UnliftIO.Resource (MonadResource, ReleaseKey)
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.NamedType ((:::))
import Vulkan.Zero (Zero(..))

import Engine.SpirV.Reflect (Reflect)
import Engine.Vulkan.DescSets (Bound(..), Compatible)
import Engine.Vulkan.Format (HasVkFormat(..), formatSize)
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, DsLayoutBindings, getPipelineCache)
import Render.Code (Code)
import Resource.Vulkan.DescriptorLayout qualified as Layout
import Resource.Vulkan.Named qualified as Named

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 [DsLayoutBindings]
  , 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 $
    Pipeline.allocateWith $ create extent msaa renderpass config

create
  :: ( MonadVulkan env io
     , HasRenderPass renderpass
     , Shader.Specialization spec
     , HasCallStack
     )
  => Maybe Vk.Extent2D
  -> Vk.SampleCountFlagBits
  -> renderpass
  -> Config dsl vertices instances spec
  -> io (Pipeline dsl vertices instances)
create mextent msaa renderpass Config{..} = withFrozenCallStack do
  -- TODO: get from outside
  dsLayouts <- Layout.create $ Vector.fromList (unTagged cDescLayouts)

  -- TODO: get from outside ?
  pipelineLayout <- Layout.forPipeline
    dsLayouts
    cPushConstantRanges
  Named.objectOrigin pipelineLayout

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

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

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

    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
  :: forall a pipeLayout vertices instances
  .  ( a ~ Pipeline pipeLayout vertices instances
     , HasVertexInputBindings vertices -- XXX: 0-2 of {positions, attrs} (e.g. position + uv)
     , HasVertexInputBindings instances -- XXX: 0+ of instance attrs (e.g. static params + dynamic transforms)
     )
  => SomeStruct Vk.PipelineVertexInputStateCreateInfo
vertexInput = 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

    bindings =
      filter (not . null . snd) $
        vertexInputBindings @vertices <>
        vertexInputBindings @instances

-- * 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
            )

type VertexInputBinding = (Vk.VertexInputRate, [Vk.Format])

vertexFormat :: forall a . HasVkFormat a => VertexInputBinding
vertexFormat = (Vk.VERTEX_INPUT_RATE_VERTEX, getVkFormat @a)

instanceFormat :: forall a . HasVkFormat a => VertexInputBinding
instanceFormat = (Vk.VERTEX_INPUT_RATE_INSTANCE, getVkFormat @a)

class HasVertexInputBindings a where
  vertexInputBindings :: [VertexInputBinding]

instance HasVertexInputBindings () where
  vertexInputBindings = []

instance HasVertexInputBindings Transform where
  vertexInputBindings = [instanceFormat @Transform]