-- {-# 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 (Word, IORef EntityPointer)
pointers :: IOVec (Word, 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#
g #)) IORef EntityPointer
pointer Entities
entities = do
  let !i' :: Int
i' = Int# -> Int
I# (Int# -> Int) -> Int# -> Int
forall a b. (a -> b) -> a -> b
$ Word# -> Int#
word2Int# Word#
i
  Entities -> Int -> IO ()
fillVec Entities
entities Int
i'
  Vec (PrimState IO) (Word, IORef EntityPointer)
-> Int -> (Word, IORef EntityPointer) -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
Vec (PrimState m) a -> Int -> a -> m ()
Vec.write Entities
entities.pointers Int
i' (Word# -> Word
W# Word#
g, 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#
g #)) Entities
entities = do
  let !i :: Int
i = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
i')
  !l <- Vec (PrimState IO) (Word, 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
      (!g', !p) <- Vec.unsafeRead entities.pointers i
      pure $ if p /= entities.nullPtr && g' == W# g 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, _) <- 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.pushBack entities.pointers 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 = !counter, free = !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) (Word, 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 (0, entities.nullPtr)

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

fillVec :: Entities -> Int -> IO ()
fillVec :: Entities -> Int -> IO ()
fillVec !Entities
entities !Int
n = do
  !l <- Vec (PrimState IO) (Word, IORef EntityPointer) -> IO Int
forall (m :: * -> *) a. PrimMonad m => Vec (PrimState m) a -> m Int
Vec.length Entities
entities.pointers
  when (n >= l) $ do
    replicateM_ (n - l + 1) $ Vec.pushBack entities.pointers (0, entities.nullPtr)

-- | Create a new storage for entities.
emptyEntities :: IO Entities
emptyEntities :: IO Entities
emptyEntities = do
  !map <- Int -> IO (Vec (PrimState IO) (Word, 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 (0, nullPtr)
  pure $ 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#
g #)) !Entities
entities = do
  let !i :: Int
i = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
i')
  !l <- Vec (PrimState IO) (Word, 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
      (!g', !p) <- Vec.unsafeRead entities.pointers i
      pure $ p /= entities.nullPtr && g' == W# g
    else pure False