Tables.hs

{-# LANGUAGE AllowAmbiguousTypes #-}

module Mischief.ECS.Tables where

import Control.Monad (forM, when)
import Data.Bifunctor qualified
import Data.Foldable (Foldable (toList), for_)
import Data.IORef
import Data.List (transpose)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (catMaybes)
import Data.Traversable (for)
import Data.Vector qualified as Vector
import GHC.Base (Int (..), eqWord#, isTrue#)
import Mischief.ECS.Components
import Mischief.ECS.Components.Bundle
import Mischief.ECS.Entities
import Mischief.ECS.Vec (IOVec)
import Mischief.ECS.Vec qualified as Vec

newtype Tables = Tables {inner :: IOVec Table}

data Table = Table
  { columns :: IORef (Map ComponentId Column),
    components :: [ComponentId],
    ticks :: IORef (Map ComponentId ComponentTick),
    entities :: IOVec (Entity, IORef EntityPointer)
  }

newtype Column = Column (IOVec ErasedComponent)

baseline_entities_cap :: Int
baseline_entities_cap = 64

emptyTables :: IO Tables
emptyTables = do
  v <- Vec.new 128
  Vec.pushBack v =<< newTable []
  pure $ Tables v

newTable :: [ComponentId] -> IO Table
newTable components = do
  cols <-
    for components $ \x -> do
      empty_vec <- Vec.new baseline_entities_cap
      pure (x, Column empty_vec)

  columns <- newIORef $ Map.fromList (toList cols)
  entities <- Vec.new baseline_entities_cap
  ticks <- newIORef $ Map.fromList (map (,ComponentTick (Tick 0)) components)

  return $ Table {columns = columns, components, entities, ticks}

tableIsEmpty :: Table -> IO Bool
tableIsEmpty table = Vec.null table.entities

-- removeTable :: ArchetypeId -> Tables -> IO ()
-- removeTable archetype tables = do
--   let Tables tables' = tables
--   modifyIORef' tables' $ Map.filterWithKey (\archetype' _ -> archetype' /= archetype)

-- |
--  runs the specified monadic action with the value of the given map key as input, if the key exists.
--  if it doesn't, do nothing
tapMap ::
  (Monad m, Ord k) =>
  Map k a ->
  k ->
  (a -> m ()) ->
  m ()
tapMap map k act = do
  maybe (pure ()) act $ Map.lookup k map

insertComponentsIntoMap :: ProcessedBundleData -> Map ComponentId Column -> IO ()
insertComponentsIntoMap bundle map =
  for_ bundle.elements $ \el ->
    tapMap map el.id $ \(Column col) ->
      Vec.pushBack col el.component

insertComponentsIntoTable :: ProcessedBundleData -> Tick -> ChangedComponents -> Table -> IO ()
insertComponentsIntoTable bundle tick changed table = do
  for_ changed $ \a -> modifyIORef' table.ticks (Map.insert a (ComponentTick tick))

  cols <- readIORef table.columns
  insertComponentsIntoMap bundle cols

replaceComponentsIntoMap ::
  ProcessedBundleData ->
  EntityPointer ->
  Map ComponentId Column ->
  IO ()
replaceComponentsIntoMap bundle (EntityPointer (# _, rowIndex #)) tableMap = do
  for_ bundle.elements $ \el ->
    tapMap tableMap el.id $ \(Column col) ->
      Vec.unsafeWrite col (I# rowIndex) el.component

--   foldl' modifyMap tableMap bundle.elements
--  where
--   modifyMap tableMap element = Map.adjust (modifyComponents element) element.id tableMap
--   modifyComponents element (Column x) =
--     let (x', _ : ys) = splitAt pointer.rowIndex x
--      in Column $ x' ++ [element.component] ++ ys

replaceComponentsIntoTable ::
  ProcessedBundleData ->
  -- | The current tick that will be as the components' change tick.
  Maybe Tick ->
  EntityPointer ->
  Table ->
  IO ()
replaceComponentsIntoTable bundle tick pointer table = do
  case tick of
    Nothing -> pure ()
    Just tick -> do
      for_ bundle.elements $ \a -> modifyIORef' table.ticks (Map.insert a.id (ComponentTick tick))

  cols <- readIORef table.columns
  replaceComponentsIntoMap bundle pointer cols

takeFromColumn :: EntityPointer -> Column -> IO ErasedComponent
takeFromColumn (EntityPointer (# _, rowIndex #)) (Column col) = Vec.takeSwap col (I# rowIndex)

takeComponentsFromTable :: EntityPointer -> Table -> IO ProcessedBundleData
takeComponentsFromTable (EntityPointer (# archetypeId, rowIndex #)) table = do
  cols <- readIORef table.columns
  newColumns <- for cols $ takeFromColumn (EntityPointer (# archetypeId, rowIndex #))
  removeEntityFromTable (I# rowIndex) table
  let elements = map (uncurry ProcessedBundleElement) $ Map.toList newColumns
  pure ProcessedBundleData {elements}

collectComponentIdsFromTable :: Table -> IO [ComponentId]
collectComponentIdsFromTable table = do
  cols <- readIORef table.columns
  return $ map fst $ Map.toList cols

-- let elements = map getComponent newColumns
-- return ProcessedBundleData{elements}

removeComponentFromColumn :: EntityPointer -> Column -> IO ()
removeComponentFromColumn (EntityPointer (# _, rowIndex #)) (Column col) = Vec.removeSwap col (I# rowIndex)

removeComponentsFromMap :: EntityPointer -> Map ComponentId Column -> IO ()
removeComponentsFromMap pointer columnMap =
  for_ columnMap $ removeComponentFromColumn pointer

-- map (\(id, column) -> (id, removeComponentFromColumn pointer column)) (Map.toList columnMap)

removeComponentsFromTable :: EntityPointer -> Table -> IO ()
removeComponentsFromTable (EntityPointer (# archetypeId, rowIndex #)) table = do
  cols <- readIORef table.columns
  removeComponentsFromMap (EntityPointer (# archetypeId, rowIndex #)) cols
  removeEntityFromTable (I# rowIndex) table

removeRow :: Int -> IOVec (Entity, IORef EntityPointer) -> IO ()
removeRow row_idx vec = do
  len <- Vec.length vec
  Vec.removeSwap vec row_idx
  when (row_idx < len - 1) $ do
    Vec.tap vec row_idx $ \(_, ptr) ->
      modifyIORef' ptr $ \(EntityPointer (# archetypeId, _ #)) ->
        let !(I# id) = row_idx
         in EntityPointer (# archetypeId, id #)

removeEntityFromTable :: Int -> Table -> IO ()
removeEntityFromTable row table = removeRow row table.entities

-- insertResourceIntoTables :: ProcessedBundleData -> Tick -> Tables -> ArchetypeId -> (Entity, IORef EntityPointer) -> IO ()
-- insertResourceIntoTables bundle tick (Tables tables) archetype (entity, pointerRef) =
--   do
--     innerTables <- readIORef tables

--     table <- newTable bundle
--     let newTables = Map.insert archetype table innerTables
--     writeIORef tables newTables

--     rowIndex <- Vec.length table.entities
--     Vec.pushBack table.entities (entity, pointerRef)

--     let !(I# archetype') = archetype.id
--     let !(I# rowIndex') = rowIndex
--     writeIORef pointerRef $ EntityPointer (# archetype', rowIndex' #)

--     insertComponentsIntoTable bundle table

type ChangedComponents = [ComponentId]

type AddedComponents = [ComponentId]

insertEntityIntoTables :: ProcessedBundleData -> Tick -> ChangedComponents -> Tables -> ArchetypeId -> (Entity, IORef EntityPointer) -> IO ()
insertEntityIntoTables bundle tick changed (Tables tables) archetype pointerRef =
  do
    table <- Vec.unsafeRead tables archetype.id

    rowIndex <- Vec.length table.entities
    Vec.pushBack table.entities pointerRef

    let !(I# archetype') = archetype.id
    let !(I# rowIndex') = rowIndex
    writeIORef (snd pointerRef) $ EntityPointer (# archetype', rowIndex' #)

    insertComponentsIntoTable bundle tick changed table

tryGetComponentFromColumn :: forall c. (Component c) => Column -> EntityPointer -> IO (Maybe c)
tryGetComponentFromColumn (Column components) (EntityPointer (# _, rowIndex #)) = do
  element <- Vec.unsafeRead components (I# rowIndex)
  pure $ tryGetComponent element

tryGetRelCollectionFromTable :: forall c. (Component c) => Table -> Entity -> EntityPointer -> ComponentId -> IO (Maybe [(Entity, Rel c)])
tryGetRelCollectionFromTable table entity pointer (ComponentId (# id, _ #)) =
  do
    columns <- readIORef table.columns
    -- TODO: improve lookup performance for partial tuples

    components' <- forM (Map.toList columns) $ \(ComponentId (# id', target' #), column) -> do
      if isTrue# $ eqWord# id id'
        then do
          case target' of
            Nothing -> return Nothing
            Just entity -> do
              component <- tryGetComponentFromColumn @c column pointer
              return $ fmap (,entity) component
        else return Nothing

    if null $ catMaybes components'
      then
        return Nothing
      else do
        return $ Just $ map (\(value, target) -> (entity, Rel value target)) $ catMaybes components'

tryGetComponentFromTable :: forall c. (Component c) => Table -> EntityPointer -> ComponentId -> IO (Maybe c)
tryGetComponentFromTable table pointer componentId =
  do
    columns <- readIORef table.columns
    let column = Map.lookup componentId columns
    case column of
      Nothing -> return Nothing
      Just column -> tryGetComponentFromColumn column pointer

tryGetRelCollectionFromTables :: forall c. (Component c) => Tables -> Entity -> EntityPointer -> ComponentId -> IO (Maybe [(Entity, Rel c)])
tryGetRelCollectionFromTables (Tables tables) entity (EntityPointer (# archetypeId, rowIndex #)) componentId =
  do
    table <- Vec.unsafeRead tables (I# archetypeId)
    tryGetRelCollectionFromTable table entity (EntityPointer (# archetypeId, rowIndex #)) componentId

tryGetComponentFromTables :: forall c. (Component c) => Tables -> EntityPointer -> ComponentId -> IO (Maybe c)
tryGetComponentFromTables (Tables tables) (EntityPointer (# archetypeId, rowIndex #)) componentId =
  do
    table <- Vec.unsafeRead tables (I# archetypeId)
    tryGetComponentFromTable table (EntityPointer (# archetypeId, rowIndex #)) componentId

-- tryGetEntityTicksFromColumn :: Column -> EntityPointer -> IO ComponentTicks
-- tryGetEntityTicksFromColumn (Column components) (EntityPointer (# _, rowIndex #)) = (\x -> x.ticks) <$> Vec.unsafeRead components (I# rowIndex)

tryGetTicksFromTable :: Table -> ComponentId -> IO (Maybe ComponentTick)
tryGetTicksFromTable table componentId = do
  ticks <- readIORef table.ticks
  pure $ Map.lookup componentId ticks

-- do
--   columns <- readIORef table.columns
--   case Map.lookup componentId columns of
--     Nothing -> do
--       l <- Vec.length table.entities
--       return $ map (const Nothing) [1 .. l]
--     Just column -> do
--       map Just <$> tryGetTicksFromColumn column

-- tryGetEntityTicksFromTable :: Table -> EntityPointer -> ComponentId -> IO (Maybe ComponentTicks)
-- tryGetEntityTicksFromTable table pointer componentId =
--   do
--     columns <- readIORef table.columns
--     case Map.lookup componentId columns of
--       Nothing -> return Nothing
--       Just column -> Just <$> tryGetEntityTicksFromColumn column pointer

tryGetTicksFromArchetype :: ArchetypeId -> IOVec Table -> ComponentId -> IO (Maybe ComponentTick)
tryGetTicksFromArchetype archetype tables componentId = do
  table <- Vec.unsafeRead tables archetype.id
  tryGetTicksFromTable table componentId

-- tryGetEntityTicksFromArchetype :: ArchetypeId -> IOVec Table -> EntityPointer -> ComponentId -> IO (Maybe ComponentTicks)
-- tryGetEntityTicksFromArchetype archetype tables pointer componentId = do
--   table <- Vec.unsafeRead tables archetype.id
--   tryGetEntityTicksFromTable table pointer componentId

tryGetTicksFromTables :: Tables -> [ArchetypeId] -> ComponentId -> IO [Maybe ComponentTick]
tryGetTicksFromTables (Tables tables) archetypes componentId =
  mapM (\archetype -> tryGetTicksFromArchetype archetype tables componentId) archetypes

-- tryGetEntityTicksFromTables :: Tables -> EntityPointer -> ComponentId -> IO (Maybe ComponentTicks)
-- tryGetEntityTicksFromTables (Tables tables) (EntityPointer (# archetypeId, rowIndex #)) componentId =
--   do
--     tryGetEntityTicksFromArchetype (ArchetypeId $ I# archetypeId) tables (EntityPointer (# archetypeId, rowIndex #)) componentId

tryGetComponentsFromColumn :: forall c. (Component c) => Column -> IO [c]
tryGetComponentsFromColumn (Column components) = do
  frozen <- Vec.freeze components
  let x = Vector.mapM tryGetComponent frozen
  pure $ maybe [] Vector.toList x

tryGetRelCollectionsFromTable :: forall c. (Component c) => Table -> ComponentId -> IO [(Entity, [Rel c])]
tryGetRelCollectionsFromTable table (ComponentId (# id, _ #)) =
  do
    -- let Just entity = componentId.entity
    columns <- readIORef table.columns
    -- TODO: improve lookup performance for partial tuples

    components' :: [Maybe [(c, Entity)]] <- forM (Map.toList columns) $ \(ComponentId (# id', target' #), column) -> do
      if isTrue# $ eqWord# id id'
        then do
          case target' of
            Nothing -> return Nothing
            Just entity -> do
              components <- tryGetComponentsFromColumn @c column
              return $ Just $ map (,entity) components
        else return Nothing

    entities <- Vec.toList table.entities
    let components'' = zip (map fst entities) $ transpose $ catMaybes components'
    return $
      map
        (Data.Bifunctor.second (map (uncurry Rel)))
        components''

tryGetComponentsFromTable :: forall c. (Component c) => Table -> ComponentId -> IO [(Entity, c)]
tryGetComponentsFromTable table componentId =
  do
    columns <- readIORef table.columns
    case Map.lookup componentId columns of
      Nothing -> undefined
      Just column -> do
        results <- tryGetComponentsFromColumn @c column
        entities <- Vec.toList table.entities
        let zipped = zip (map fst entities) results
        return $ map (\(e, r) -> (e, r)) zipped

tryGetEntitiesFromTable :: Table -> IO [Entity]
tryGetEntitiesFromTable table =
  do
    entities <- Vec.toList table.entities
    return $ map fst entities

tryGetComponentsFromTableMaybe :: forall c. (Component c) => Table -> ComponentId -> IO [(Entity, Maybe c)]
tryGetComponentsFromTableMaybe table componentId =
  do
    columns <- readIORef table.columns
    case Map.lookup componentId columns of
      Nothing -> do
        entities <- Vec.toList table.entities
        return $ map ((,Nothing) . fst) entities
      Just column -> do
        results <- tryGetComponentsFromColumn @c column
        entities <- Vec.toList table.entities
        return $ zipWith (\x e -> (e, Just x)) results (map fst entities)

tryGetRelCollectionsFromArchetype :: forall c. (Component c) => ArchetypeId -> IOVec Table -> ComponentId -> IO [(Entity, [Rel c])]
tryGetRelCollectionsFromArchetype archetype tables componentId = do
  table <- Vec.unsafeRead tables archetype.id
  tryGetRelCollectionsFromTable table componentId

tryGetComponentsFromArchetype :: forall c. (Component c) => ArchetypeId -> IOVec Table -> ComponentId -> IO [(Entity, c)]
tryGetComponentsFromArchetype archetype tables componentId = do
  table <- Vec.unsafeRead tables archetype.id
  tryGetComponentsFromTable table componentId

tryGetEntitiesFromArchetype :: ArchetypeId -> IOVec Table -> IO [Entity]
tryGetEntitiesFromArchetype archetype tables = do
  table <- Vec.unsafeRead tables archetype.id
  tryGetEntitiesFromTable table

tryGetComponentsFromArchetypeMaybe :: forall c. (Component c) => ArchetypeId -> IOVec Table -> ComponentId -> IO [(Entity, Maybe c)]
tryGetComponentsFromArchetypeMaybe archetype tables componentId = do
  table <- Vec.unsafeRead tables archetype.id
  tryGetComponentsFromTableMaybe table componentId

tryGetRelCollectionsFromTables :: forall c. (Component c) => Tables -> [ArchetypeId] -> ComponentId -> IO [(Entity, [Rel c])]
tryGetRelCollectionsFromTables (Tables tables) archetypes componentId =
  do
    results <- mapM (\archetype -> tryGetRelCollectionsFromArchetype archetype tables componentId) archetypes
    return $ concat results

tryGetComponentsFromTables :: forall c. (Component c) => Tables -> [ArchetypeId] -> ComponentId -> IO [(Entity, c)]
tryGetComponentsFromTables (Tables tables) archetypes componentId =
  do
    results <- mapM (\archetype -> tryGetComponentsFromArchetype archetype tables componentId) archetypes
    return $ concat results

tryGetEntitiesFromTables :: Tables -> [ArchetypeId] -> IO [Entity]
tryGetEntitiesFromTables (Tables tables) archetypes =
  do
    results <- mapM (`tryGetEntitiesFromArchetype` tables) archetypes
    return $ concat results

tryGetComponentsFromTablesMaybe :: forall c. (Component c) => Tables -> [ArchetypeId] -> ComponentId -> IO [(Entity, Maybe c)]
tryGetComponentsFromTablesMaybe (Tables tables) archetypes componentId =
  do
    results <- mapM (\archetype -> tryGetComponentsFromArchetypeMaybe archetype tables componentId) archetypes
    return $ concat results

-- newtype Result c = Result (c, Entity)

-- data ErasedResult where
--   ErasedResult :: Result c -> ErasedResult

-- value :: Result c -> c
-- value (Result (c, _)) = c

-- type family IsComp a where
--   IsComp (Rel a) = False
--   IsComp a = True

-- entityOf :: Result c -> Entity
-- entityOf (Result (_, e)) = e

-- instance (Show c) => Show (Result c) where
--   show :: Result c -> String
--   show = show . value

-- instance (Eq c) => Eq (Result c) where
--   (==) :: Result c -> Result c -> Bool
--   (==) a b = value a == value b

-- instance (Ord c) => Ord (Result c) where
--   compare :: Result c -> Result c -> Ordering
--   compare a b = compare (value a) (value b)

-- instance (HasField a b c) => HasField a (Result b) c where
--   getField a = getField @a (value a)

-- class DeepValue' flag c i | flag c -> i where
--   deepValue' :: c -> i

-- instance DeepValue' True (Result c) c where
--   deepValue' = value

-- instance DeepValue' False (Result (Rel c)) c where
--   deepValue' x = x.comp

-- class DeepValue c i | c -> i where
--   deepValue :: c -> i

-- instance (DeepValue' (IsComp c) (Result c) i) => DeepValue (Result c) i where
--   deepValue = deepValue' @(IsComp c)