{-# OPTIONS_GHC -fplugin Foreign.Storable.Generic.Plugin #-}

module Resource.Mesh.Types
  ( AxisAligned(..)

  , Meta(..)

  , NodeGroup(..)
  , NodePartitions(..)

  , Nodes
  , Node(..)
  , TexturedNodes
  , TexturedNode(..)
  , TextureParams(..) -- XXX: copypasta from LitTextured

  , Measurements(..)
  , measureAa
  , measureAaWith
  , middle
  , middleAa
  , size
  , sizeAa

  , HasRange(..)

  , encodeStorable
  , decodeStorable
  ) where

import RIO

import Codec.Serialise qualified as CBOR
import Codec.Serialise.Decoding qualified as CBOR (Decoder, decodeBytes)
import Codec.Serialise.Encoding qualified as CBOR (Encoding)
import Control.Foldl qualified as L
import Data.ByteString.Unsafe qualified as BS
import Data.Typeable (typeRep, typeRepTyCon)
import Foreign (Storable(..), castPtr)
import Foreign qualified
import Foreign.Storable.Generic (GStorable)
import Geomancy (Transform(..), Vec2, Vec4, withVec3)
import Geomancy.Mat4 qualified as Mat4
import Geomancy.Vec3 qualified as Vec3
import RIO.Vector.Storable qualified as Storable
import System.IO.Unsafe (unsafePerformIO)
import Vulkan.Zero (Zero(..))

import Resource.Model (IndexRange(..))

data AxisAligned a = AxisAligned
  { aaX :: a
  , aaY :: a
  , aaZ :: a
  }
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable, Generic)

instance Applicative AxisAligned where
  pure x = AxisAligned
    { aaX = x
    , aaY = x
    , aaZ = x
    }

  funcs <*> args = AxisAligned
    { aaX = aaX funcs $ aaX args
    , aaY = aaY funcs $ aaY args
    , aaZ = aaZ funcs $ aaZ args
    }

instance Storable a => Storable (AxisAligned a) where
  alignment ~_ = alignment (error "AxisAligned.alignment" :: a)

  sizeOf ~_ = 3 * sizeOf (error "AxisAligned.sizeOf" :: a)

  peek ptr = do
    aaX <- peekElemOff (castPtr ptr) 0
    aaY <- peekElemOff (castPtr ptr) 1
    aaZ <- peekElemOff (castPtr ptr) 2
    pure AxisAligned{..}

  poke ptr AxisAligned{..} = do
    pokeElemOff (castPtr ptr) 0 aaX
    pokeElemOff (castPtr ptr) 1 aaY
    pokeElemOff (castPtr ptr) 2 aaZ

instance (CBOR.Serialise a) => CBOR.Serialise (AxisAligned a)

-- * Whole-scene metadata

data Meta = Meta
  { -- XXX: full-scene draws
    mOpaqueIndices  :: IndexRange
  , mBlendedIndices :: IndexRange

    -- XXX: per-node draws
  , mOpaqueNodes    :: IndexRange
  , mBlendedNodes   :: IndexRange

  , mBoundingSphere :: Vec4
  , mTransformBB    :: Transform
  , mMeasurements   :: AxisAligned Measurements
  }
  deriving (Show, Generic)

instance GStorable Meta

instance Eq Meta where
  a == b = and
    [ mBoundingSphere a == mBoundingSphere b
    , mOpaqueIndices  a == mOpaqueIndices  b
    , mBlendedIndices a == mBlendedIndices b
    , mOpaqueNodes    a == mOpaqueNodes    b
    , mBlendedNodes   a == mBlendedNodes   b
    , mMeasurements   a == mMeasurements   b

    , Mat4.toListRowMajor (mTransformBB a) ==
      Mat4.toListRowMajor (mTransformBB b)
    ]

instance CBOR.Serialise Meta where
  encode Meta{..} = mconcat
    [ CBOR.encode    mOpaqueIndices
    , CBOR.encode    mBlendedIndices
    , CBOR.encode    mOpaqueNodes
    , CBOR.encode    mBlendedNodes
    , encodeStorable mBoundingSphere
    , encodeStorable mTransformBB
    , CBOR.encode    mMeasurements
    ]

  decode = do
    mOpaqueIndices  <- CBOR.decode
    mBlendedIndices <- CBOR.decode
    mOpaqueNodes    <- CBOR.decode
    mBlendedNodes   <- CBOR.decode
    mBoundingSphere <- decodeStorable
    mTransformBB    <- decodeStorable
    mMeasurements   <- CBOR.decode
    pure Meta{..}

-- * Scene parts

data NodeGroup
  = NodeOpaque
  | NodeBlended
  -- TODO: NodeCutout
  deriving (Eq, Ord, Show, Enum, Bounded)

data NodePartitions a = NodePartitions
  { npOpaque  :: a
  , npBlended :: a
  }
  deriving (Eq, Show, Functor, Foldable, Traversable)

type Nodes = Storable.Vector Node

data Node = Node
  { nBoundingSphere :: Vec4
  , nTransformBB    :: Transform
  , nRange          :: IndexRange
  , nMeasurements   :: AxisAligned Measurements
  }
  deriving (Show, Generic)

instance Eq Node where
  a == b = and
    [ nBoundingSphere a == nBoundingSphere b
    , nMeasurements   a == nMeasurements   b
    , nRange          a == nRange          b

    , Mat4.toListRowMajor (nTransformBB a) ==
      Mat4.toListRowMajor (nTransformBB b)
    ]

instance GStorable Node

type TexturedNodes = Storable.Vector TexturedNode

data TexturedNode = TexturedNode
  { tnNode              :: Node
  , tnBase              :: TextureParams
  , tnEmissive          :: TextureParams
  , tnNormal            :: TextureParams
  , tnOcclusion         :: TextureParams
  , tnMetallicRoughness :: TextureParams
  }
  deriving (Eq, Show, Generic)

instance GStorable TexturedNode

-- XXX: copypasta from LitTextured.Model
data TextureParams = TextureParams
  { tpScale     :: Vec2
  , tpOffset    :: Vec2
  , tpGamma     :: Vec4
  , tpSamplerId :: Int32
  , tpTextureId :: Int32
  }
  deriving (Eq, Show, Generic)

instance Zero TextureParams where
  zero = TextureParams
    { tpScale     = 1
    , tpOffset    = 0
    , tpGamma     = 1.0
    , tpSamplerId = minBound
    , tpTextureId = minBound
    }

instance Storable TextureParams where
  alignment ~_ = 4

  sizeOf ~_ = 8 + 8 + 16 + 4 + 4

  poke ptr TextureParams{..} = do
    pokeByteOff ptr  0 tpScale
    pokeByteOff ptr  8 tpOffset
    pokeByteOff ptr 16 tpGamma
    pokeByteOff ptr 32 tpSamplerId
    pokeByteOff ptr 36 tpTextureId

  peek ptr = do
    tpScale     <- peekByteOff ptr  0
    tpOffset    <- peekByteOff ptr  8
    tpGamma     <- peekByteOff ptr 16
    tpSamplerId <- peekByteOff ptr 32
    tpTextureId <- peekByteOff ptr 36
    pure TextureParams{..}

-- * Measurements

data Measurements = Measurements
  { mMin  :: Float
  , mMax  :: Float
  , mMean :: Float
  , mStd  :: Float
  }
  deriving (Eq, Ord, Show, Generic)

instance Storable Measurements where
  alignment ~_ = 4 -- XXX: 16?

  sizeOf ~_ = 4 * 4

  peek ptr = do
    mMin  <- peekByteOff ptr  0
    mMax  <- peekByteOff ptr  4
    mMean <- peekByteOff ptr  8
    mStd  <- peekByteOff ptr 12
    pure Measurements{..}

  poke ptr Measurements{..} = do
    pokeByteOff ptr  0 mMin
    pokeByteOff ptr  4 mMax
    pokeByteOff ptr  8 mMean
    pokeByteOff ptr 12 mStd

instance CBOR.Serialise Measurements

{-# INLINEABLE middleAa #-}
middleAa :: AxisAligned Measurements -> AxisAligned Float
middleAa = fmap middle

{-# INLINEABLE middle #-}
middle :: Measurements -> Float
middle Measurements{mMax, mMin} = mMin * 0.5 + mMax * 0.5

{-# INLINEABLE sizeAa #-}
sizeAa :: AxisAligned Measurements -> AxisAligned Float
sizeAa = fmap size

{-# INLINEABLE size #-}
size :: Measurements -> Float
size Measurements{mMax, mMin} = mMax - mMin

measureAaWith
  :: (Foldable outer, Foldable inner)
  => (a -> inner Vec3.Packed)
  -> outer a
  -> AxisAligned Measurements
measureAaWith f = L.fold (measureAaWithF f)

measureAaWithF :: (Foldable t) => (a -> (t Vec3.Packed)) -> L.Fold a (AxisAligned Measurements)
measureAaWithF f = L.premap f (L.handles L.folded measureAaF)

measureAa :: Foldable t => t Vec3.Packed -> AxisAligned Measurements
measureAa = L.fold measureAaF

measureAaF :: L.Fold Vec3.Packed (AxisAligned Measurements)
measureAaF = AxisAligned
  <$> L.premap packedX measureF
  <*> L.premap packedY measureF
  <*> L.premap packedZ measureF

measureF :: L.Fold Float Measurements
measureF = do
  mMin  <- fmap (fromMaybe 0) L.minimum
  mMax  <- fmap (fromMaybe 0) L.maximum
  mMean <- L.mean
  mStd  <- L.std
  pure Measurements{..}

-- * Utils

{-# INLINE packedX #-}
packedX :: Vec3.Packed -> Float
packedX (Vec3.Packed pos) = withVec3 pos \x _y _z -> x

{-# INLINE packedY #-}
packedY :: Vec3.Packed -> Float
packedY (Vec3.Packed pos) = withVec3 pos \_x y _z -> y

{-# INLINE packedZ #-}
packedZ :: Vec3.Packed -> Float
packedZ (Vec3.Packed pos) = withVec3 pos \_x _y z -> z

class HasRange a where
  getRange :: a -> IndexRange
  adjustRange :: a -> Word32 -> a

instance HasRange Node where
  {-# INLINEABLE getRange #-}
  getRange = nRange

  {-# INLINEABLE adjustRange #-}
  adjustRange node@Node{nRange} newFirstIndex = node
    { nRange = nRange
        { irFirstIndex = newFirstIndex
        }
    }

instance HasRange TexturedNode where
  {-# INLINE getRange #-}
  getRange = getRange . tnNode

  {-# INLINE adjustRange #-}
  adjustRange tn@TexturedNode{tnNode} newFirstIndex = tn
    { tnNode = adjustRange tnNode newFirstIndex
    }

-- | CBOR.encode helper for storable types (vectors, etc.)
{-# NOINLINE encodeStorable #-}
encodeStorable :: forall a . Storable a => a -> CBOR.Encoding
encodeStorable x = unsafePerformIO do
  ptr <- Foreign.malloc @a
  poke ptr x
  buf <- BS.unsafePackMallocCStringLen
    ( Foreign.castPtr ptr
    , Foreign.sizeOf (undefined :: a)
    )
  pure $ CBOR.encode buf

-- | CBOR.decode helper for storable types (vectors, etc.)
decodeStorable :: forall a s . (Storable a, Typeable a) => CBOR.Decoder s a
decodeStorable = do
  buf <- CBOR.decodeBytes
  case fromBuf buf of
    Left err ->
      fail err
    Right !res ->
      pure res
  where
    expected =
      Foreign.sizeOf (undefined :: a)

    {-# NOINLINE fromBuf #-}
    fromBuf :: ByteString -> Either String a
    fromBuf buf =
      unsafePerformIO $
        BS.unsafeUseAsCStringLen buf \(ptr, len) ->
          if len == expected then do
            !res <- Foreign.peek (Foreign.castPtr ptr)
            pure $ Right res
          else
            pure . Left $ mconcat
              [ "Storable size mismatch for "
              , show (typeRepTyCon . typeRep $ Proxy @a)
              , " (expected: ", show expected
              , ", got: ", show len
              , ")"
              ]