ClassAttributeMap.hs

module Rattletrap.Type.ClassAttributeMap where

import qualified Data.IntMap as IntMap
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.Text as Text
import qualified Data.Tuple as Tuple
import qualified Rattletrap.Data as Data
import qualified Rattletrap.Type.AttributeMapping as AttributeMapping
import qualified Rattletrap.Type.Cache as Cache
import qualified Rattletrap.Type.ClassMapping as ClassMapping
import qualified Rattletrap.Type.CompressedWord as CompressedWord
import qualified Rattletrap.Type.List as RList
import qualified Rattletrap.Type.Str as Str
import qualified Rattletrap.Type.U32 as U32

-- | This data structure holds all the information about classes, objects, and
-- attributes in the replay. The class hierarchy is not fixed; it is encoded
-- in the 'Rattletrap.Content.Content'. Similarly, the attributes that belong
-- to each class are not fixed either. Converting the raw data into a usable
-- structure is tedious; see 'make'.
data ClassAttributeMap = ClassAttributeMap
  { -- | A map from object IDs to their names.
    objectMap :: Map.Map U32.U32 Str.Str,
    -- | A map from object IDs to their class IDs.
    objectClassMap :: Map.Map U32.U32 U32.U32,
    -- | A map from class IDs to a map from attribute stream IDs to attribute
    -- IDs.
    value :: Map.Map U32.U32 (Map.Map U32.U32 U32.U32),
    nameMap :: IntMap.IntMap Str.Str
  }
  deriving (Eq, Show)

type Bimap l r = (Map.Map l r, Map.Map r l)

bimap :: (Ord l, Ord r) => [(l, r)] -> Bimap l r
bimap xs = (Map.fromList xs, Map.fromList (fmap Tuple.swap xs))

lookupL :: (Ord l) => l -> Bimap l r -> Maybe r
lookupL k = Map.lookup k . fst

lookupR :: (Ord r) => r -> Bimap l r -> Maybe l
lookupR k = Map.lookup k . snd

-- | Makes a 'ClassAttributeMap' given the necessary fields from the
-- 'Rattletrap.Content.Content'.
make ::
  -- | From 'Rattletrap.Content.objects'.
  RList.List Str.Str ->
  -- | From 'Rattletrap.Content.classMappings'.
  RList.List ClassMapping.ClassMapping ->
  -- | From 'Rattletrap.Content.caches'.
  RList.List Cache.Cache ->
  -- | From 'Rattletrap.Content.names'.
  RList.List Str.Str ->
  ClassAttributeMap
make objects classMappings caches names =
  let objectMap_ = makeObjectMap objects
      classMap = makeClassMap classMappings
      objectClassMap_ = makeObjectClassMap objectMap_ classMap
      classCache = makeClassCache classMap caches
      attributeMap = makeAttributeMap caches
      classIds = fmap (\(_, classId, _, _) -> classId) classCache
      parentMap = makeParentMap classCache
      value_ =
        Map.fromList
          ( fmap
              ( \classId ->
                  let ownAttributes =
                        Maybe.fromMaybe Map.empty (Map.lookup classId attributeMap)
                      parentsAttributes = case Map.lookup classId parentMap of
                        Nothing -> []
                        Just parentClassIds ->
                          fmap
                            ( \parentClassId ->
                                Maybe.fromMaybe
                                  Map.empty
                                  (Map.lookup parentClassId attributeMap)
                            )
                            parentClassIds
                      attributes = ownAttributes : parentsAttributes
                   in (classId, Map.fromList (concatMap Map.toList attributes))
              )
              classIds
          )
      nameMap_ = makeNameMap names
   in ClassAttributeMap objectMap_ objectClassMap_ value_ nameMap_

makeNameMap :: RList.List Str.Str -> IntMap.IntMap Str.Str
makeNameMap names =
  IntMap.fromDistinctAscList (zip [0 ..] (RList.toList names))

getName :: IntMap.IntMap Str.Str -> U32.U32 -> Maybe Str.Str
getName nameMap_ nameIndex =
  IntMap.lookup (fromIntegral (U32.toWord32 nameIndex)) nameMap_

makeObjectClassMap ::
  Map.Map U32.U32 Str.Str ->
  Bimap U32.U32 Str.Str ->
  Map.Map U32.U32 U32.U32
makeObjectClassMap objectMap_ classMap = do
  let objectIds = Map.keys objectMap_
  let classIds = fmap (getClassId objectMap_ classMap) objectIds
  let rawPairs = zip objectIds classIds
  let pairs =
        Maybe.mapMaybe
          ( \(objectId, maybeClassId) -> case maybeClassId of
              Nothing -> Nothing
              Just classId -> Just (objectId, classId)
          )
          rawPairs
  Map.fromList pairs

getClassId ::
  Map.Map U32.U32 Str.Str ->
  Bimap U32.U32 Str.Str ->
  U32.U32 ->
  Maybe U32.U32
getClassId objectMap_ classMap objectId = do
  objectName <- getObjectName objectMap_ objectId
  className <- getClassName objectName
  lookupR className classMap

makeClassCache ::
  Bimap U32.U32 Str.Str ->
  RList.List Cache.Cache ->
  [(Maybe Str.Str, U32.U32, U32.U32, U32.U32)]
makeClassCache classMap caches =
  fmap
    ( \cache ->
        let classId = Cache.classId cache
         in ( lookupL classId classMap,
              classId,
              Cache.cacheId cache,
              Cache.parentCacheId cache
            )
    )
    (RList.toList caches)

makeClassMap :: RList.List ClassMapping.ClassMapping -> Bimap U32.U32 Str.Str
makeClassMap classMappings =
  bimap
    ( fmap
        ( \classMapping ->
            (ClassMapping.streamId classMapping, ClassMapping.name classMapping)
        )
        (RList.toList classMappings)
    )

makeAttributeMap ::
  RList.List Cache.Cache -> Map.Map U32.U32 (Map.Map U32.U32 U32.U32)
makeAttributeMap caches =
  Map.fromList
    ( fmap
        ( \cache ->
            ( Cache.classId cache,
              Map.fromList
                ( fmap
                    ( \attributeMapping ->
                        ( AttributeMapping.streamId attributeMapping,
                          AttributeMapping.objectId attributeMapping
                        )
                    )
                    (RList.toList (Cache.attributeMappings cache))
                )
            )
        )
        (RList.toList caches)
    )

makeShallowParentMap ::
  [(Maybe Str.Str, U32.U32, U32.U32, U32.U32)] -> Map.Map U32.U32 U32.U32
makeShallowParentMap classCache =
  Map.fromList
    ( Maybe.mapMaybe
        ( \xs -> case xs of
            [] -> Nothing
            (maybeClassName, classId, _, parentCacheId) : rest -> do
              parentClassId <- getParentClass maybeClassName parentCacheId rest
              pure (classId, parentClassId)
        )
        (List.tails (reverse classCache))
    )

makeParentMap ::
  [(Maybe Str.Str, U32.U32, U32.U32, U32.U32)] -> Map.Map U32.U32 [U32.U32]
makeParentMap classCache =
  let shallowParentMap = makeShallowParentMap classCache
   in Map.mapWithKey
        (\classId _ -> getParentClasses shallowParentMap classId)
        shallowParentMap

getParentClasses :: Map.Map U32.U32 U32.U32 -> U32.U32 -> [U32.U32]
getParentClasses shallowParentMap classId =
  case Map.lookup classId shallowParentMap of
    Nothing -> []
    Just parentClassId ->
      parentClassId : getParentClasses shallowParentMap parentClassId

getParentClass ::
  Maybe Str.Str ->
  U32.U32 ->
  [(Maybe Str.Str, U32.U32, U32.U32, U32.U32)] ->
  Maybe U32.U32
getParentClass maybeClassName parentCacheId xs = case maybeClassName of
  Nothing -> getParentClassById parentCacheId xs
  Just className -> getParentClassByName className parentCacheId xs

getParentClassById ::
  U32.U32 -> [(Maybe Str.Str, U32.U32, U32.U32, U32.U32)] -> Maybe U32.U32
getParentClassById parentCacheId xs =
  case dropWhile (\(_, _, cacheId, _) -> cacheId /= parentCacheId) xs of
    [] ->
      if parentCacheId == U32.fromWord32 0
        then Nothing
        else
          getParentClassById
            (U32.fromWord32 (U32.toWord32 parentCacheId - 1))
            xs
    (_, parentClassId, _, _) : _ -> Just parentClassId

getParentClassByName ::
  Str.Str ->
  U32.U32 ->
  [(Maybe Str.Str, U32.U32, U32.U32, U32.U32)] ->
  Maybe U32.U32
getParentClassByName className parentCacheId xs =
  case Map.lookup (Str.toText className) Data.parentClasses of
    Nothing -> getParentClassById parentCacheId xs
    Just parentClassName ->
      Maybe.maybe
        (getParentClassById parentCacheId xs)
        Just
        ( Maybe.listToMaybe
            ( fmap
                (\(_, parentClassId, _, _) -> parentClassId)
                ( filter
                    (\(_, _, cacheId, _) -> cacheId <= parentCacheId)
                    ( filter
                        ( \(maybeClassName, _, _, _) ->
                            fmap Str.toText maybeClassName == Just parentClassName
                        )
                        xs
                    )
                )
            )
        )

makeObjectMap :: RList.List Str.Str -> Map.Map U32.U32 Str.Str
makeObjectMap objects =
  Map.fromAscList (zip (fmap U32.fromWord32 [0 ..]) (RList.toList objects))

getObjectName :: Map.Map U32.U32 Str.Str -> U32.U32 -> Maybe Str.Str
getObjectName objectMap_ objectId = Map.lookup objectId objectMap_

getClassName :: Str.Str -> Maybe Str.Str
getClassName rawObjectName =
  fmap Str.fromText $
    Map.lookup
      (Str.toText $ normalizeObjectName rawObjectName)
      Data.objectClasses

normalizeObjectName :: Str.Str -> Str.Str
normalizeObjectName objectName =
  let name = Str.toText objectName
      crowdActor = Text.pack "TheWorld:PersistentLevel.CrowdActor_TA"
      crowdManager = Text.pack "TheWorld:PersistentLevel.CrowdManager_TA"
      boostPickup = Text.pack "TheWorld:PersistentLevel.VehiclePickup_Boost_TA"
      mapScoreboard = Text.pack "TheWorld:PersistentLevel.InMapScoreboard_TA"
      breakout = Text.pack "TheWorld:PersistentLevel.BreakOutActor_Platform_TA"
   in if Text.isInfixOf crowdActor name
        then Str.fromText crowdActor
        else
          if Text.isInfixOf crowdManager name
            then Str.fromText crowdManager
            else
              if Text.isInfixOf boostPickup name
                then Str.fromText boostPickup
                else
                  if Text.isInfixOf mapScoreboard name
                    then Str.fromText mapScoreboard
                    else
                      if Text.isInfixOf breakout name
                        then Str.fromText breakout
                        else objectName

classHasLocation :: Str.Str -> Bool
classHasLocation className =
  Set.member (Str.toText className) Data.classesWithLocation

classHasRotation :: Str.Str -> Bool
classHasRotation className =
  Set.member (Str.toText className) Data.classesWithRotation

getAttributeIdLimit :: Map.Map U32.U32 U32.U32 -> Maybe Word
getAttributeIdLimit attributeMap = do
  ((streamId, _), _) <- Map.maxViewWithKey attributeMap
  pure (fromIntegral (U32.toWord32 streamId))

getAttributeName ::
  ClassAttributeMap ->
  Map.Map U32.U32 U32.U32 ->
  CompressedWord.CompressedWord ->
  Maybe Str.Str
getAttributeName classAttributeMap attributeMap streamId = do
  let key = U32.fromWord32 (fromIntegral (CompressedWord.value streamId))
  attributeId <- Map.lookup key attributeMap
  let objectMap_ = objectMap classAttributeMap
  Map.lookup attributeId objectMap_

getAttributeMap ::
  ClassAttributeMap ->
  Map.Map CompressedWord.CompressedWord U32.U32 ->
  CompressedWord.CompressedWord ->
  Maybe (Map.Map U32.U32 U32.U32)
getAttributeMap classAttributeMap actorMap actorId = do
  objectId <- Map.lookup actorId actorMap
  let objectClassMap_ = objectClassMap classAttributeMap
  classId <- Map.lookup objectId objectClassMap_
  let value_ = value classAttributeMap
  Map.lookup classId value_