module Render.Lit.Material.Collect
  ( LoadedModel
  , SceneModel(..)
  , sceneMaterials
  , modelMaterials
  , nodeMaterials
  ) where

import RIO

import RIO.Map qualified as Map
import RIO.Set qualified as Set
import RIO.Vector.Storable qualified as Storable

import Render.Lit.Material (shiftTextures)
import Render.Lit.Material.Model qualified as LitMaterial
import Resource.Buffer qualified as Buffer
import Resource.Mesh.Types qualified as Mesh
import Resource.Mesh.Lit qualified as Lit
import Render.Lit.Material (Material)

type LoadedModel =
  ( Mesh.Meta
  , Storable.Vector Lit.MaterialNode
  , LitMaterial.Model 'Buffer.Staged
  )

data SceneModel models textures = SceneModel
  { forall models textures. SceneModel models textures -> Text
smLabel            :: Text
  , forall models textures.
SceneModel models textures -> models -> LoadedModel
smGetModel         :: models -> LoadedModel
  , forall models textures.
SceneModel models textures -> textures -> Int32
smGetTextureOffset :: textures -> Int32
  }

sceneMaterials
  :: Foldable t
  => models
  -> textures
  -> t (SceneModel models textures)
  -> Storable.Vector Material
sceneMaterials :: forall (t :: * -> *) models textures.
Foldable t =>
models
-> textures -> t (SceneModel models textures) -> Vector Material
sceneMaterials models
loadedModels textures
combinedTextures =
  forall a. Storable a => [a] -> Vector a
Storable.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Map Text (Set Int), Map Int (Text, Material))
-> SceneModel models textures
-> (Map Text (Set Int), Map Int (Text, Material))
beep (forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty)
  where
    beep :: (Map Text (Set Int), Map Int (Text, Material))
-> SceneModel models textures
-> (Map Text (Set Int), Map Int (Text, Material))
beep (Map Text (Set Int), Map Int (Text, Material))
acc SceneModel{Text
models -> LoadedModel
textures -> Int32
smGetTextureOffset :: textures -> Int32
smGetModel :: models -> LoadedModel
smLabel :: Text
$sel:smGetTextureOffset:SceneModel :: forall models textures.
SceneModel models textures -> textures -> Int32
$sel:smGetModel:SceneModel :: forall models textures.
SceneModel models textures -> models -> LoadedModel
$sel:smLabel:SceneModel :: forall models textures. SceneModel models textures -> Text
..} =
      let
        (Meta
_meta, Vector MaterialNode
materialNodes, Model 'Staged
_model) = models -> LoadedModel
smGetModel models
loadedModels
        textureOffset :: Int32
textureOffset = textures -> Int32
smGetTextureOffset textures
combinedTextures
      in
        forall (t :: * -> *).
Foldable t =>
Text
-> Int32
-> (Map Text (Set Int), Map Int (Text, Material))
-> t MaterialNode
-> (Map Text (Set Int), Map Int (Text, Material))
modelMaterials Text
smLabel Int32
textureOffset (Map Text (Set Int), Map Int (Text, Material))
acc (forall a. Storable a => Vector a -> [a]
Storable.toList Vector MaterialNode
materialNodes)

modelMaterials
  :: Foldable t
  => Text
  -> Int32
  -> (Map Text (Set Int), Map Int (Text, Material))
  -> t Lit.MaterialNode
  -> (Map Text (Set Int), Map Int (Text, Material))
modelMaterials :: forall (t :: * -> *).
Foldable t =>
Text
-> Int32
-> (Map Text (Set Int), Map Int (Text, Material))
-> t MaterialNode
-> (Map Text (Set Int), Map Int (Text, Material))
modelMaterials Text
label Int32
textureOffset =
  forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Text
-> Int32
-> MaterialNode
-> (Map Text (Set Int), Map Int (Text, Material))
-> (Map Text (Set Int), Map Int (Text, Material))
nodeMaterials Text
label Int32
textureOffset)

nodeMaterials
  :: Text
  -> Int32
  -> Lit.MaterialNode
  -> (Map Text (Set Int), Map Int (Text, Material))
  -> (Map Text (Set Int), Map Int (Text, Material))
nodeMaterials :: Text
-> Int32
-> MaterialNode
-> (Map Text (Set Int), Map Int (Text, Material))
-> (Map Text (Set Int), Map Int (Text, Material))
nodeMaterials Text
label Int32
textureStart Lit.MaterialNode{Int
$sel:mnMaterialIx:MaterialNode :: MaterialNode -> Int
mnMaterialIx :: Int
mnMaterialIx, Material
$sel:mnMaterial:MaterialNode :: MaterialNode -> Material
mnMaterial :: Material
mnMaterial} (Map Text (Set Int)
ids, Map Int (Text, Material)
collection) =
  case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int
mnMaterialIx Map Int (Text, Material)
collection of
    Maybe (Text, Material)
Nothing ->
      ( forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Monoid a => a -> a -> a
mappend Text
label (forall a. a -> Set a
Set.singleton Int
mnMaterialIx) Map Text (Set Int)
ids
      , forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Int
mnMaterialIx (Text
label, Material
newMaterial) Map Int (Text, Material)
collection
      )
    Just (Text
oldLabel, Material
oldMaterial) ->
      if Material
newMaterial forall a. Eq a => a -> a -> Bool
/= Material
oldMaterial then
        forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
          [ String
"Ouf... The material indices are clashing for " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Text
label
          , String
"ID: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Int
mnMaterialIx forall a. Semigroup a => a -> a -> a
<> String
" from " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Text
oldLabel
          , String
""
          , String
"Old: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Material
oldMaterial
          , String
"Now: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Material
mnMaterial
          , String
""
          , String
"Bad material offset? Known material IDs:"
          , [String] -> String
unlines do
              (Text
knownLabel, Set Int
knownIds) <- forall k a. Map k a -> [(k, a)]
Map.toList Map Text (Set Int)
ids
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show Text
knownLabel forall a. Semigroup a => a -> a -> a
<> String
": " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall a. Set a -> [a]
Set.toList Set Int
knownIds)
          , String
"Try setting " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Text
label forall a. Semigroup a => a -> a -> a
<> String
" material offset to " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Int
suggestOffset
          ]
      else
        (Map Text (Set Int)
ids, Map Int (Text, Material)
collection)
  where
    newMaterial :: Material
newMaterial =
      Int32 -> Material -> Material
shiftTextures Int32
textureStart Material
mnMaterial

    suggestOffset :: Int
    suggestOffset :: Int
suggestOffset =
      (forall a. Num a => a -> a -> a
+ Int
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. Ord a => a -> a -> a
max Int
0 forall a b. (a -> b) -> a -> b
$
        forall a. [Maybe a] -> [a]
catMaybes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Set a -> Maybe a
Set.lookupMax forall a b. (a -> b) -> a -> b
$
          forall k a. Map k a -> [a]
Map.elems (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Text
label Map Text (Set Int)
ids)