-- {-# 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 GHC.Base
import GHC.Conc
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