-- {-# LANGUAGE DeriveLift #-}

module Mischief.ECS.Entities
  ( -- * Entity
    Entity (..),
    EntityPointer (..),
    getPointer,
    isAliveIO,

    -- * Storage
    Entities,
    EntityCounter,
    getNewEntity,
    removeEntity,
    insertPointer,
    getNewEntityComp,
    emptyEntities,
  )
where

import Control.Concurrent.STM.TVar
import Control.Monad
import Data.IORef
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (isJust)
import GHC.Base
import GHC.Conc
import Mischief.ECS.Components
import Mischief.ECS.EntityDef
import Mischief.ECS.Vec (IOVec)
import Mischief.ECS.Vec qualified as Vec

-- | A pointer to the exact table and row that an entity is in.
--
-- The first @Int#@ is the id of its archetype. The second is the index of the row it's in.
data EntityPointer = EntityPointer (# Int#, Int# #)

-- { archetypeId :: ArchetypeId,
--   rowIndex :: Int
-- }
-- deriving (Show)

-- | A storage for entity ids and pointers.
data Entities = Entities
  { -- | Associates each @Entity@ to an @EntityPointer@.
    Entities -> Vec RealWorld (IORef EntityPointer)
pointers :: IOVec (IORef EntityPointer),
    Entities -> IORef EntityPointer
nullPtr :: IORef EntityPointer,
    -- | A counter for assigning new entity ids. It is in a TVar so
    -- it can be used in parallel systems.
    Entities -> TVar EntityCounter
counter :: TVar EntityCounter
  }

-- | A counter and a list for recycling entities.
data EntityCounter = EntityCounter {EntityCounter -> Word#
counter :: Word#, EntityCounter -> [Entity]
free :: [Entity]}

-- | Associate an Entity with a new pointer.
insertPointer :: Entity -> IORef EntityPointer -> Entities -> IO ()
insertPointer :: Entity -> IORef EntityPointer -> Entities -> IO ()
insertPointer (Entity (# Word#
i, Word#
_ #)) IORef EntityPointer
pointer Entities
entities = do
  Vec (PrimState IO) (IORef EntityPointer)
-> Int -> IORef EntityPointer -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> a -> m ()
Vec.write Entities
entities.pointers (Int# -> Int
I# (Int# -> Int) -> Int# -> Int
forall a b. (a -> b) -> a -> b
$ Word# -> Int#
word2Int# Word#
i) IORef EntityPointer
pointer

-- | Get the pointer to an entity.
getPointer :: Entity -> Entities -> IO (Maybe (IORef EntityPointer))
getPointer :: Entity -> Entities -> IO (Maybe (IORef EntityPointer))
getPointer (Entity (# Word#
i', Word#
_ #)) Entities
entities = do
  let i :: Int
i = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
i')
  l <- Vec (PrimState IO) (IORef EntityPointer) -> IO Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
Vec.length Entities
entities.pointers

  if i < l
    then do
      p <- Vec.unsafeRead entities.pointers i
      pure $ if p /= entities.nullPtr then Just p else Nothing
    else
      pure Nothing

-- | Create a new entity.
getNewEntity :: Entities -> IO Entity
getNewEntity :: Entities -> IO Entity
getNewEntity Entities
entities = do
  (e, reused) <- STM (Entity, Bool) -> IO (Entity, Bool)
forall a. STM a -> IO a
atomically (STM (Entity, Bool) -> IO (Entity, Bool))
-> STM (Entity, Bool) -> IO (Entity, Bool)
forall a b. (a -> b) -> a -> b
$ do
    EntityCounter {counter, free} <- TVar EntityCounter -> STM EntityCounter
forall a. TVar a -> STM a
readTVar Entities
entities.counter
    case free of
      [] -> do
        TVar EntityCounter -> EntityCounter -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar Entities
entities.counter EntityCounter {counter :: Word#
counter = Word# -> Word# -> Word#
plusWord# Word#
counter Word#
1##, [Entity]
free :: [Entity]
free :: [Entity]
free}
        (Entity, Bool) -> STM (Entity, Bool)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((# Word#, Word# #) -> Entity
Entity (# Word#
counter, Word#
1## #), Bool
False)
      (Entity (# Word#
id, Word#
gen #) : [Entity]
xs) -> do
        TVar EntityCounter -> EntityCounter -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar Entities
entities.counter EntityCounter {counter :: Word#
counter = Word#
counter, free :: [Entity]
free = [Entity]
xs}
        (Entity, Bool) -> STM (Entity, Bool)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((# Word#, Word# #) -> Entity
Entity (# Word#
id, Word# -> Word# -> Word#
plusWord# Word#
gen Word#
1## #), Bool
True)

  unless reused $ do
    Vec.pushBack entities.pointers entities.nullPtr

  pure e

-- | Creates a new entity using a fresh id, never incrementing the generation of a previous one.
getNewEntityComp :: Entities -> IO Entity
getNewEntityComp :: Entities -> IO Entity
getNewEntityComp Entities
entities = do
  Vec (PrimState IO) (IORef EntityPointer)
-> IORef EntityPointer -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> a -> m ()
Vec.pushBack Entities
entities.pointers Entities
entities.nullPtr
  STM Entity -> IO Entity
forall a. STM a -> IO a
atomically (STM Entity -> IO Entity) -> STM Entity -> IO Entity
forall a b. (a -> b) -> a -> b
$ do
    EntityCounter {counter, free} <- TVar EntityCounter -> STM EntityCounter
forall a. TVar a -> STM a
readTVar Entities
entities.counter
    let e = (# Word#, Word# #) -> Entity
Entity (# Word#
counter, Word#
0## #)
    writeTVar entities.counter EntityCounter {counter = plusWord# counter 1##, free}
    pure e

-- | Remove an entity from storage.
removeEntity :: Entity -> Entities -> IO ()
removeEntity :: Entity -> Entities -> IO ()
removeEntity (Entity (# Word#
i', Word#
g #)) Entities
entities = do
  let i :: Int
i = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
i')
  l <- Vec (PrimState IO) (IORef EntityPointer) -> IO Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
Vec.length Entities
entities.pointers

  when (i < l) $ do
    Vec.unsafeWrite entities.pointers i entities.nullPtr

    atomically $ do
      EntityCounter {counter, free} <- readTVar entities.counter
      writeTVar entities.counter EntityCounter {counter, free = Entity (# i', g #) : free}

-- | Create a new storage for entities.
emptyEntities :: IO Entities
emptyEntities :: IO Entities
emptyEntities = do
  map <- Int -> IO (Vec (PrimState IO) (IORef EntityPointer))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (Vec (PrimState m) a)
Vec.new Int
256
  counter <- newTVarIO EntityCounter {counter = 1##, free = []}
  nullPtr <- newIORef $ EntityPointer (# 0#, 0# #)

  Vec.pushBack map nullPtr
  return $ Entities map nullPtr counter

-- | Check if an Entity is alive through IO.
isAliveIO :: Entity -> Entities -> IO Bool
isAliveIO :: Entity -> Entities -> IO Bool
isAliveIO (Entity (# Word#
i', Word#
_ #)) Entities
entities = do
  let i :: Int
i = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
i')
  l <- Vec (PrimState IO) (IORef EntityPointer) -> IO Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
Vec.length Entities
entities.pointers

  if i < l
    then do
      p <- Vec.unsafeRead entities.pointers i
      pure $ p /= entities.nullPtr
    else pure False