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
  { smLabel            :: Text
  , smGetModel         :: models -> LoadedModel
  , smGetTextureOffset :: textures -> Int32
  }

sceneMaterials
  :: Foldable t
  => models
  -> textures
  -> t (SceneModel models textures)
  -> Storable.Vector Material
sceneMaterials loadedModels combinedTextures =
  Storable.fromList . map snd . Map.elems . snd . foldl' beep (mempty, mempty)
  where
    beep acc SceneModel{..} =
      let
        (_meta, materialNodes, _model) = smGetModel loadedModels
        textureOffset = smGetTextureOffset combinedTextures
      in
        modelMaterials smLabel textureOffset acc (Storable.toList 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 label textureOffset =
  foldr (nodeMaterials label textureOffset)

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

    suggestOffset :: Int
    suggestOffset =
      (+ 1) . foldr max 0 $
        catMaybes . map Set.lookupMax $
          Map.elems (Map.delete label ids)