module Resource.Gltf.Scene where

import RIO

import Codec.GlTF qualified as GlTF
import Codec.GlTF.Mesh qualified as GlTF (MeshIx(..))
import Codec.GlTF.Node qualified as GlTF (Node, NodeIx(..))
import Codec.GlTF.Node qualified as Node
import Codec.GlTF.Root qualified as Root
import Codec.GlTF.Scene qualified as Scene
import Data.Coerce (coerce)
import Data.Tree (Tree)
import Data.Tree qualified as Tree
import Geomancy (Transform(..), quaternion)
import Geomancy.Mat4 qualified as Mat4
import Geomancy.Transform qualified as Transform
import Geomancy.Vec3 qualified as Vec3
import RIO.Vector qualified as Vector

import Resource.Gltf.Model (Mesh, MeshPrimitive, Stuff(..), VertexAttrs(..))

unfoldSceneM
  :: HasLogFunc env
  => Int
  -> Transform
  -> Root.GlTF
  -> Vector Mesh
  -> RIO env (Tree SceneNode)
unfoldSceneM materialOffset initialTransform root allMeshes = do
  (allNodes, initialNodes) <- getRootNodes root
  meshes <- unfoldNodesM allMeshes allNodes (toList initialNodes)
  pure Tree.Node
    { rootLabel = SceneNode
        { snOrigin     = 0
        , snPrimitives = Nothing
        , snNode       = emptyNode
        }
    , subForest =
        map (toSceneNode materialOffset initialTransform) meshes
    }

{-# INLINEABLE toSceneNode #-}
toSceneNode
  :: Int
  -> Transform
  -> Tree (Maybe (Vector MeshPrimitive), GlTF.Node)
  -> Tree SceneNode
toSceneNode materialOffset initialTransform =
  fmap (injectTransforms materialOffset) .
  collectTransforms initialTransform

getRootNodes :: HasLogFunc env => GlTF.GlTF -> RIO env (Vector Node.Node, Vector Node.NodeIx)
getRootNodes root = do
  rootNodes <- case Root.scenes root of
    Nothing -> do
      logWarn "No scenes"
      pure Nothing
    Just scenes ->
      case Vector.toList scenes of
        [] -> do
          logWarn "Empty scene vector"
          pure Nothing
        [one] ->
          pure $ Scene.nodes one
        pick : _rest -> do
          logWarn $ mconcat
            [ "Picking first scene among "
            , display (Vector.length scenes)
            ]
          pure $ Scene.nodes pick

  allNodes <- case Root.nodes root of
    Nothing ->
      throwString "TODO: fallback to raw meshes"
    Just nodes -> do
      pure nodes

  case rootNodes of
    Nothing ->
      throwString "TODO: fallback for lack of scene"
    Just start -> do
      pure (allNodes, start)

emptyNode :: Node.Node
emptyNode = Node.Node
  { camera      = Nothing
  , children    = Nothing
  , skin        = Nothing
  , matrix      = Nothing
  , mesh        = Nothing
  , rotation    = Nothing
  , scale       = Nothing
  , translation = Nothing
  , weights     = Nothing
  , name        = Nothing
  , extensions  = Nothing
  , extras      = Nothing
  }

data LookupError
  = NodeNotFound Int
  | MeshNotFound Int
  deriving (Eq, Ord, Show)

instance Exception LookupError

injectTransforms
  :: Int
  -> ( Maybe (Vector MeshPrimitive)
     , Transform
     , GlTF.Node
     )
  -> SceneNode
injectTransforms materialOffset (mmesh, transform, snNode) = SceneNode{..}
  where
    snOrigin = Vec3.Packed $ Transform.apply 0 transform

    snPrimitives = fmap (Vector.map $ bimap adjustMaterial adjustNode) mmesh

    adjustMaterial = fmap \(materialId, gltfMaterial) ->
      ( materialId + materialOffset
      , gltfMaterial
      )

    adjustNode Stuff{..} = Stuff
      { sPositions = fmap applyTransform sPositions
      , sAttrs     = fmap applyAttrTransform sAttrs
      , sIndices   = sIndices
      }
      where
        applyTransform pos =
          coerce $ Transform.apply (coerce pos) transform

        applyTransformDir dir =
          coerce $ Vec3.normalize $ Transform.apply (coerce dir) (Transform transformDir)
          where
          transformDir = Mat4.transpose . Mat4.inverse $
            Mat4.pointwise (unTransform transform) nullifyTranslation (*)

        applyAttrTransform va = va
          { vaNormal  = applyTransformDir (vaNormal va)
          , vaTangent = applyTransformDir (vaTangent va)
          }

collectTransforms
  :: Transform
  -> Tree (Maybe (Vector MeshPrimitive), GlTF.Node)
  -> Tree (Maybe (Vector MeshPrimitive), Transform, GlTF.Node)
collectTransforms initial root = Tree.unfoldTree go (initial, root)
  where
    go (parent, Tree.Node{rootLabel=(mmesh, node), subForest}) =
      let
        collected = localTransform node <> parent
      in
        ( (mmesh, collected, node)
        , map (collected,) subForest
        )

-- | Build node tree and shed lookup errors as exception.
unfoldNodesM
  :: MonadThrow m
  => Vector Mesh
  -> Vector GlTF.Node
  -> [GlTF.NodeIx]
  -> m [Tree (Maybe (Vector MeshPrimitive), GlTF.Node)]
unfoldNodesM allMeshes allNodes =
  either throwM pure . traverse (inflateNode allMeshes allNodes)

-- | Combine lookup operations for nodes and meshes.
inflateNode
  :: Vector Mesh
  -> Vector GlTF.Node
  -> GlTF.NodeIx
  -> Either LookupError (Tree (Maybe (Vector MeshPrimitive), GlTF.Node))
inflateNode allMeshes allNodes startNode =
  for (unfoldNode allNodes startNode) \getNode -> do
    node <- getNode
    mesh <- getMesh allMeshes node
    pure (mesh, node)

unfoldNode
  :: Vector GlTF.Node
  -> GlTF.NodeIx
  -> Tree (Either LookupError GlTF.Node)
unfoldNode allNodes = Tree.unfoldTree fetch
  where
    fetch (GlTF.NodeIx ix) =
      case allNodes Vector.!? ix of
        Nothing ->
          ( Left $ NodeNotFound ix
          , []
          )
        Just node ->
          ( Right node
          , maybe [] toList $ Node.children node
          )

getMesh
  :: Vector mesh
  -> GlTF.Node
  -> Either LookupError (Maybe mesh)
getMesh allMeshes node =
  case Node.mesh node of
    Nothing ->
      Right Nothing
    Just (GlTF.MeshIx ix) ->
      case allMeshes Vector.!? ix of
        Nothing ->
          Left $ MeshNotFound ix
        Just mesh ->
          Right (Just mesh)

localTransform :: Node.Node -> Transform
localTransform node = mconcat
  [ nodeMatrix
  , nodeScale, nodeRotate, nodeTranslate
  ]
  where
    nodeScale = case Node.scale node of
      Nothing ->
        mempty
      Just (sx, sy, sz) ->
        Transform.scale3 sx sy sz

    nodeTranslate =
      case Node.translation node of
        Nothing ->
          mempty
        Just (tx, ty, tz) ->
          Transform.translate tx ty tz

    nodeRotate = case Node.rotation node of
      Nothing ->
        mempty
      Just (qx, qy, qz, qW) ->
        Transform.rotateQ $ quaternion qW qx qy qz

    nodeMatrix = case Node.matrix node of
      Nothing ->
        mempty
      Just (Node.NodeMatrix memoryBytes) ->
        case Vector.toList memoryBytes of
          [ x1, x2, x3, x4,
            y1, y2, y3, y4,
            z1, z2, z3, z4,
            w1, w2, w3, w4 ] ->
              Mat4.colMajor @Transform
                x1 y1 z1 w1
                x2 y2 z2 w2
                x3 y3 z3 w3
                x4 y4 z4 w4
          _ ->
            error "Node matrix isn't 16-element"

nullifyTranslation :: Mat4.Mat4
nullifyTranslation =
  Mat4.rowMajor
    1 1 1 0
    1 1 1 0
    1 1 1 0
    0 0 0 1

data SceneNode = SceneNode
  { snOrigin     :: ~Vec3.Packed
  , snPrimitives :: Maybe (Vector MeshPrimitive)
  , snNode       :: GlTF.Node
  }
  deriving (Show)