{-# 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(..)
  ) where

import RIO

import Control.Foldl qualified as L
import Foreign (Storable(..), castPtr)
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 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

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

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

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