module Rattletrap.ClassAttributeMap where

import Rattletrap.ActorMap
import Rattletrap.AttributeMapping
import Rattletrap.Cache
import Rattletrap.ClassMapping
import Rattletrap.Data
import Rattletrap.Primitive

import qualified Data.Bimap as Bimap
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.Text as Text

-- | 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 'makeClassAttributeMap'.
data ClassAttributeMap = ClassAttributeMap
  { classAttributeMapObjectMap :: Map.Map Word32 Text
  -- ^ A map from object IDs to their names.
  , classAttributeMapObjectClassMap :: Map.Map Word32 Word32
  -- ^ A map from object IDs to their class IDs.
  , classAttributeMapValue :: Map.Map Word32 (Map.Map Word32 Word32)
  -- ^ A map from class IDs to a map from attribute stream IDs to attribute
  -- IDs.
  , classAttributeMapNameMap :: IntMap.IntMap Text
  } deriving (Eq, Ord, Show)

-- | Makes a 'ClassAttributeMap' given the necessary fields from the
-- 'Rattletrap.Content.Content'.
makeClassAttributeMap
  :: List Text
  -- ^ From 'Rattletrap.Content.contentObjects'.
  -> List ClassMapping
  -- ^ From 'Rattletrap.Content.contentClassMappings'.
  -> List Cache
  -- ^ From 'Rattletrap.Content.contentCaches'.
  -> List Text
  -- ^ From 'Rattletrap.Content.contentNames'.
  -> ClassAttributeMap
makeClassAttributeMap objects classMappings caches names =
  let
    objectMap = makeObjectMap objects
    classMap = makeClassMap classMappings
    objectClassMap = makeObjectClassMap objectMap classMap
    classCache = makeClassCache classMap caches
    attributeMap = makeAttributeMap caches
    classIds = map (\(_, classId, _, _) -> classId) classCache
    parentMap = makeParentMap classCache
    value = Map.fromList
      ( map
        ( \classId ->
          let
            ownAttributes =
              Maybe.fromMaybe Map.empty (Map.lookup classId attributeMap)
            parentsAttributes = case Map.lookup classId parentMap of
              Nothing -> []
              Just parentClassIds -> map
                ( \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 :: List Text -> IntMap.IntMap Text
makeNameMap names = IntMap.fromDistinctAscList (zip [0 ..] (listValue names))

getName :: IntMap.IntMap Text -> Word32 -> Maybe Text
getName nameMap nameIndex =
  IntMap.lookup (fromIntegral (word32Value nameIndex)) nameMap

makeObjectClassMap
  :: Map.Map Word32 Text -> Bimap.Bimap Word32 Text -> Map.Map Word32 Word32
makeObjectClassMap objectMap classMap = do
  let objectIds = Map.keys objectMap
  let classIds = map (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 Word32 Text -> Bimap.Bimap Word32 Text -> Word32 -> Maybe Word32
getClassId objectMap classMap objectId = do
  objectName <- getObjectName objectMap objectId
  className <- getClassName objectName
  Bimap.lookupR className classMap

makeClassCache
  :: Bimap.Bimap Word32 Text
  -> List Cache
  -> [(Maybe Text, Word32, Word32, Word32)]
makeClassCache classMap caches = map
  ( \cache ->
    let
      classId = cacheClassId cache
    in
      ( Bimap.lookup classId classMap
      , classId
      , cacheCacheId cache
      , cacheParentCacheId cache
      )
  )
  (listValue caches)

makeClassMap :: List ClassMapping -> Bimap.Bimap Word32 Text
makeClassMap classMappings = Bimap.fromList
  ( map
    ( \classMapping ->
      (classMappingStreamId classMapping, classMappingName classMapping)
    )
    (listValue classMappings)
  )

makeAttributeMap :: List Cache -> Map.Map Word32 (Map.Map Word32 Word32)
makeAttributeMap caches = Map.fromList
  ( map
    ( \cache ->
      ( cacheClassId cache
      , Map.fromList
        ( map
          ( \attributeMapping ->
            ( attributeMappingStreamId attributeMapping
            , attributeMappingObjectId attributeMapping
            )
          )
          (listValue (cacheAttributeMappings cache))
        )
      )
    )
    (listValue caches)
  )

makeShallowParentMap
  :: [(Maybe Text, Word32, Word32, Word32)] -> Map.Map Word32 Word32
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 Text, Word32, Word32, Word32)] -> Map.Map Word32 [Word32]
makeParentMap classCache =
  let
    shallowParentMap = makeShallowParentMap classCache
  in
    Map.mapWithKey
      (\classId _ -> getParentClasses shallowParentMap classId)
      shallowParentMap

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

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

getParentClassById
  :: Word32 -> [(Maybe Text, Word32, Word32, Word32)] -> Maybe Word32
getParentClassById parentCacheId xs =
  case dropWhile (\(_, _, cacheId, _) -> cacheId /= parentCacheId) xs of
    [] -> if parentCacheId == Word32 0
      then Nothing
      else getParentClassById (Word32 (word32Value parentCacheId - 1)) xs
    (_, parentClassId, _, _):_ -> Just parentClassId

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

parentClasses :: Map.Map Text Text
parentClasses = Map.map
  stringToText
  (Map.mapKeys stringToText (Map.fromList rawParentClasses))

makeObjectMap :: List Text -> Map.Map Word32 Text
makeObjectMap objects =
  Map.fromAscList (zip (map Word32 [0 ..]) (listValue objects))

getObjectName :: Map.Map Word32 Text -> Word32 -> Maybe Text
getObjectName objectMap objectId = Map.lookup objectId objectMap

getClassName :: Text -> Maybe Text
getClassName rawObjectName =
  Map.lookup (normalizeObjectName rawObjectName) objectClasses

normalizeObjectName :: Text -> Text
normalizeObjectName objectName =
  let
    name = textValue 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 Text crowdActor
    else if Text.isInfixOf crowdManager name
      then Text crowdManager
      else if Text.isInfixOf boostPickup name
        then Text boostPickup
        else if Text.isInfixOf mapScoreboard name
          then Text mapScoreboard
          else if Text.isInfixOf breakout name
            then Text breakout
            else objectName

objectClasses :: Map.Map Text Text
objectClasses = Map.map
  stringToText
  (Map.mapKeys stringToText (Map.fromList rawObjectClasses))

classHasLocation :: Text -> Bool
classHasLocation className = Set.member className classesWithLocation

classesWithLocation :: Set.Set Text
classesWithLocation = Set.fromList (map stringToText rawClassesWithLocation)

classHasRotation :: Text -> Bool
classHasRotation className = Set.member className classesWithRotation

classesWithRotation :: Set.Set Text
classesWithRotation = Set.fromList (map stringToText rawClassesWithRotation)

getAttributeIdLimit :: Map.Map Word32 Word32 -> Maybe Word
getAttributeIdLimit attributeMap = do
  ((streamId, _), _) <- Map.maxViewWithKey attributeMap
  let limit = fromIntegral (word32Value streamId)
  pure limit

getAttributeName
  :: ClassAttributeMap -> Map.Map Word32 Word32 -> CompressedWord -> Maybe Text
getAttributeName classAttributeMap attributeMap streamId = do
  let key = Word32 (fromIntegral (compressedWordValue streamId))
  attributeId <- Map.lookup key attributeMap
  let objectMap = classAttributeMapObjectMap classAttributeMap
  Map.lookup attributeId objectMap

getAttributeMap
  :: ClassAttributeMap
  -> ActorMap
  -> CompressedWord
  -> Maybe (Map.Map Word32 Word32)
getAttributeMap classAttributeMap actorMap actorId = do
  objectId <- Map.lookup actorId actorMap
  let objectClassMap = classAttributeMapObjectClassMap classAttributeMap
  classId <- Map.lookup objectId objectClassMap
  let value = classAttributeMapValue classAttributeMap
  Map.lookup classId value