Asset.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedDatatypes #-}

module Mischief.Assets.Internal.Asset where

import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.State (MonadState (get, put), State, runState)
import Data.Data
import Data.Foldable (for_)
import Data.HashTable.IO qualified as H
import Data.Kind (Type)
import Data.Maybe
import Data.Set (Set)
import Data.Set qualified as Set
import Data.String (IsString)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Exts (IsString (..))
import KDL (DecodeValue (valueDecoder))
import KDL qualified
import Mischief.Assets.Internal.Path
import Mischief.ECS
import Mischief.ECS.Entities
import Mischief.ECS.Resources
import Path (Path)
import Path qualified

class (Component a) => Asset a where
  type Loadable a :: Bool
  type Loadable a = True

  type Saveable a :: Bool
  type Saveable a = True

  loadAsset :: FilePath -> Text -> IO (Either Text a)
  default loadAsset :: (Loadable a ~ False) => FilePath -> Text -> IO (Either Text a)
  loadAsset _ _ = pure $ Left $ T.show (typeRep $ Proxy @a) <> " does not implement a load function!"

  saveAsset :: a -> FilePath -> Text -> IO (Either Text ())
  default saveAsset :: (Saveable a ~ False) => a -> FilePath -> Text -> IO (Either Text ())
  saveAsset _ _ _ = pure $ Left $ T.show (typeRep $ Proxy @a) <> " does not implement a save function!"

  rec :: a -> Rec ()
  rec = const $ pure ()

data RecAsset where
  RecAsset :: AssetRef a -> RecAsset

data LoadRec = LoadRec deriving (Component)

newtype RecAssets = RecAssets [RecAsset]

type Rec = State RecAssets

addRec :: AssetRef a -> Rec ()
addRec a' = do
  RecAssets a <- get
  put $ RecAssets $ RecAsset a' : a

type RelativePath = Text

data Loading = Loading deriving (Component)

data Loaded = Loaded deriving (Component)

data AssetEntity = AssetEntity deriving (Component)

loadByRelPath :: forall a. (Asset a, Bundle a) => Text -> System Entity
loadByRelPath !path' = do
  path <- relAssetToAbsFile path'
  loadByAbsPath @a path path'

loadByAbsPath :: forall a. (Asset a, Bundle a) => Path Path.Abs Path.File -> Text -> System Entity
loadByAbsPath !path !relPath = do
  a <- tryFindAsset @a path
  relDir <- T.pack . Path.toFilePath . Path.parent <$> Path.parseRelFile (T.unpack relPath)

  case a of
    Just a -> pure a
    Nothing -> do
      entity <- spawn (AssetEntity, Loading, AssetType $ Proxy @a, AssetFilePath {absFilePath = path, relFilePath = relPath})

      Just (AssetRegistry registry) <- res @(AssetRegistry a)
      AssetTypeRegistry typeRegistry <- resOrInsertM @AssetTypeRegistry $ AssetTypeRegistry <$> liftIO H.new
      void $ liftIO $ H.mutate typeRegistry path $ \case
        Just x -> (Just $ Set.insert (AssetType $ Proxy @a) x, path)
        Nothing -> (Just $ Set.singleton (AssetType $ Proxy @a), path)

      liftIO $ H.insert registry path entity

      runAfter
        ( do
            loadAsset @a (Path.toFilePath path) relDir
        )
        ( \(!assetData) -> do
            case assetData of
              Left t -> err [i|Error when loading asset: #{t}|]
              Right assetData -> do
                Just loadRec <- single [q|entity. Has LoadRec|]

                remove (C @Loading, C @LoadRec) entity
                insert (assetData, Loaded) entity
                trigger (OnLoad @a entity)

                when loadRec $ handleLoadRec assetData (Set.singleton $ typeRep $ Proxy @a)
        )

      pure entity

handleLoadRec :: (Asset a) => a -> Set TypeRep -> System ()
handleLoadRec a prev = do
  let (_, RecAssets recs) = runState (rec a) (RecAssets [])
  for_ recs $ \(RecAsset a) -> void $ getAsset a True prev

data AssetFilePath = AssetFilePath {absFilePath :: Path Path.Abs Path.File, relFilePath :: Text} deriving (Component)

reloadByRelPath :: forall a. (Asset a, Bundle a) => Text -> System Entity
reloadByRelPath !path' = do
  path <- relAssetToAbsFile path'
  reloadByAbsPath @a path path'

reloadByAbsPath :: forall a. (Asset a, Bundle a) => Path Path.Abs Path.File -> Text -> System Entity
reloadByAbsPath !path !relPath = do
  a <- tryFindAsset @a path
  relDir <- T.pack . Path.toFilePath . Path.parent <$> Path.parseRelFile (T.unpack relPath)

  entity <- case a of
    Just a -> pure a
    Nothing -> do
      spawn (AssetEntity, Loading, AssetType $ Proxy @a, AssetFilePath {absFilePath = path, relFilePath = relPath})

  Just (AssetRegistry registry) <- res @(AssetRegistry a)
  AssetTypeRegistry typeRegistry <- resOrInsertM @AssetTypeRegistry $ AssetTypeRegistry <$> liftIO H.new
  void $ liftIO $ H.mutate typeRegistry path $ \case
    Just x -> (Just $ Set.insert (AssetType $ Proxy @a) x, path)
    Nothing -> (Just $ Set.singleton (AssetType $ Proxy @a), path)

  liftIO $ H.insert registry path entity

  runAfter
    ( do
        loadAsset @a (Path.toFilePath path) relDir
    )
    ( \(!assetData) -> do
        case assetData of
          Left t -> err [i|Error when loading asset: #{t}|]
          Right assetData -> do
            remove (C @Loading) entity
            insert (assetData, Loaded) entity
            trigger (OnLoad @a entity)
    )

  pure entity

-- retrieve' :: forall a. (Asset a) => Text -> System (Maybe a)
-- retrieve' !path' = do
--   e <- retrieveE @a path'
--   case e of
--     Nothing -> pure Nothing
--     Just e -> single [q|e. a|]

-- retrieveE :: forall a. (Asset a) => Text -> System (Maybe Entity)
-- retrieveE !path' = do
--   path <- relAssetToAbsFile path'
--   tryFindAsset @a path

store :: forall a. (Asset a, Bundle a) => Text -> a -> System Entity
store !path' !a = do
  path <- relAssetToAbsFile path'
  a' <- tryFindAsset @a path

  entity <- case a' of
    Just e -> do
      insert a e
      pure e
    Nothing -> do
      entity <- spawn $ bundle (AssetEntity, Loading, AssetType $ Proxy @a, a)
      when (T.head path' /= '!') $ insert AssetFilePath {absFilePath = path, relFilePath = path'} entity

      Just (AssetRegistry registry) <- res @(AssetRegistry a)
      liftIO $ H.insert registry path entity

      AssetTypeRegistry typeRegistry <- resOrInsertM @AssetTypeRegistry $ AssetTypeRegistry <$> liftIO H.new
      void $ liftIO $ H.mutate typeRegistry path $ \case
        Just x -> (Just $ Set.insert (AssetType $ Proxy @a) x, path)
        Nothing -> (Just $ Set.singleton (AssetType $ Proxy @a), path)

      pure entity

  trigger (OnLoad @a entity)
  pure entity

save :: forall a. (Asset a) => Entity -> System ()
save e = do
  t <- single [q|e. a, AssetFilePath|]
  case t of
    Nothing -> warn $ "Asset of type " <> T.show (typeRep $ Proxy @a) <> " not found at " <> T.show e <> "."
    Just (a, AssetFilePath {absFilePath = path, relFilePath = relPath}) -> do
      runAfter
        ( do
            saveAsset a (Path.toFilePath path) relPath
        )
        ( \case
            Left t -> warn t
            _ -> pure ()
        )

tryFindAsset :: forall a. (Asset a) => Path Path.Abs Path.File -> System (Maybe Entity)
tryFindAsset !path = do
  AssetRegistry registry <- resOrInsertM @(AssetRegistry a) (AssetRegistry <$> liftIO H.new)
  liftIO $ H.lookup registry path

-- newtype Path = Path AssetPath deriving anyclass (Component)

newtype OnLoad a = OnLoad {entity :: Entity}
  deriving anyclass (Event)
  deriving stock (Show)

data AssetType where
  AssetType :: (Asset a) => Proxy a -> AssetType
  deriving (Component)

instance Eq AssetType where
  AssetType a == AssetType b = typeRep a == typeRep b

instance Ord AssetType where
  compare (AssetType a) (AssetType b) = compare (typeRep a) (typeRep b)

type HashMap k v = H.BasicHashTable k v

newtype AssetRegistry a = AssetRegistry (HashMap (Path Path.Abs Path.File) Entity) deriving anyclass (Component)

newtype AssetTypeRegistry = AssetTypeRegistry (HashMap (Path Path.Abs Path.File) (Set AssetType)) deriving anyclass (Component)

data AssetRef a where
  Ref :: (Asset a, Bundle a) => Text -> AssetRef a
  EntityRef :: (Asset a, Bundle a) => Entity -> AssetRef a
  SubRef :: AssetRef b -> (b -> a) -> AssetRef a
  TransRef :: forall a b. AssetRef b -> (b -> AssetRef a) -> AssetRef a
  PureRef :: forall a. a -> AssetRef a
  AppRef :: AssetRef (a -> b) -> AssetRef a -> AssetRef b
  NullRef :: AssetRef a

decodeRef :: (Asset a, Bundle a) => Text -> KDL.ValueDecoder (AssetRef a)
decodeRef path = do
  t <- valueDecoder @Text
  pure $ Ref (path <> "/" <> t)

-- | Gets an AssetRef's entity. If it is unloaded, it will create an entity and return it.
getRefEntity :: forall a. AssetRef a -> System (Maybe Entity)
getRefEntity (Ref p') = do
  p <- relAssetToAbsFile p'
  e <- tryFindAsset @a p
  case e of
    Nothing -> Just <$> loadByAbsPath @a p p'
    Just x -> pure $ Just x
-- registry <- res @(AssetRegistry a)
-- case registry of
--   Nothing -> pure Nothing
--   Just (AssetRegistry registry) -> do
--     liftIO $ H.lookup registry p
getRefEntity (EntityRef e) = pure $ Just e
getRefEntity (SubRef _ _) = do
  warn "Attempted to retrieve the entity of a sub asset (data inside a loaded asset). Returning Nothing"
  pure Nothing
getRefEntity (TransRef ref f) = do
  a <- getAsset' ref
  case a of
    Nothing -> pure Nothing
    Just a -> getRefEntity $ f a
getRefEntity (PureRef _) = do
  warn "Attempted to retrieve the entity of a pure AssetRef. Returning Nothing."
  pure Nothing
getRefEntity (AppRef _ _) = do
  warn "Attempted to retrieve the entity of an AppRef. Returning Nothing."
  pure Nothing
getRefEntity NullRef = pure Nothing

-- | Gets an AssetRef's entity. Won't trigger any IO.
getRefEntity' :: forall a w m. (MonadSystem w m) => AssetRef a -> m (Maybe Entity)
getRefEntity' (Ref p) = do
  registry <- res @(AssetRegistry a)
  case registry of
    Nothing -> pure Nothing
    Just (AssetRegistry registry) -> do
      p <- relAssetToAbsFile p
      liftIO $ H.lookup registry p
getRefEntity' (EntityRef e) = pure $ Just e
getRefEntity' (SubRef _ _) = do
  warn "Attempted to retrieve the entity of a sub asset (data inside a loaded asset). Returning Nothing"
  pure Nothing
getRefEntity' (TransRef ref f) = do
  a <- getAsset' ref
  case a of
    Nothing -> pure Nothing
    Just a -> getRefEntity' $ f a
getRefEntity' (PureRef _) = do
  warn "Attempted to retrieve the entity of a pure AssetRef. Returning Nothing"
  pure Nothing
getRefEntity' (AppRef _ _) = do
  warn "Attempted to retrieve the entity of an AppRef. Returning Nothing."
  pure Nothing
getRefEntity' NullRef = pure Nothing

-- | Try to get an asset by its AssetRef. Won't trigger any IO.
getAsset' :: forall a w m. (MonadSystem w m) => AssetRef a -> m (Maybe a)
getAsset' (Ref a) = do
  e <- getRefEntity' (Ref @a a)
  case e of
    Nothing -> pure Nothing
    Just e -> single [q|e. a|]
getAsset' (EntityRef e) = single [q|e. a|]
getAsset' (SubRef a f) = fmap f <$> getAsset' a
getAsset' (TransRef a f) = do
  x <- getAsset' a
  case x of
    Nothing -> pure Nothing
    Just x -> do
      getAsset' (f x)
getAsset' (PureRef a) = pure $ Just a
getAsset' (AppRef f a) = do
  x <- getAsset' f
  y <- getAsset' a
  let z = (,) <$> x <*> y
  case z of
    Nothing -> pure Nothing
    Just (f, a) -> do
      pure $ Just $ f a
getAsset' NullRef = pure Nothing

-- loaded :: forall a w m. (MonadSystem w m) => AssetRef a -> m Bool
-- loaded ref = isJust <$> getAsset' ref

type IsRec = Bool

-- | Try to get an AssetRef's value. Will trigger it to load it it's currently unloaded.
getAsset :: forall a. AssetRef a -> IsRec -> Set TypeRep -> System (Maybe a)
getAsset (EntityRef e) b prev = do
  r <- single [q|e. Maybe a, Maybe AssetFilePath, Has Loading, Has LoadRec|]
  case r of
    Nothing -> Nothing <$ warn [i|Trying to get asset of invalid entity #{e}|]
    Just (Just a, _, _, _) -> do
      when b $ do
        let t = typeRep $ Proxy @a
        when (Set.member t prev) $ handleLoadRec a (Set.insert t prev)
      pure $ Just a
    Just (_, Nothing, _, _) -> Nothing <$ warn [i|Trying to load an asset without a file path #{e}|]
    Just (_, Just (AssetFilePath {absFilePath, relFilePath}), False, _) -> do
      void $ reloadByAbsPath @a absFilePath relFilePath
      when b $ insert LoadRec e
      pure Nothing
    Just (_, _, True, loadRec) -> do
      when (b && not loadRec) $ insert LoadRec e
      pure Nothing
getAsset (Ref path) b' prev = do
  e <- getRefEntity' (Ref @a path)
  case e of
    Just e -> do
      b <- single $ mkGet e (M @a, Has @Loading, Has @LoadRec)
      case b of
        Nothing -> Nothing <$ loadByRelPath @a path
        Just (Just a, _, _) -> do
          when b' $ do
            let t = typeRep $ Proxy @a
            when (Set.member t prev) $ handleLoadRec a (Set.insert t prev)
          pure $ Just a
        Just (Nothing, False, _) -> do
          void $ loadByRelPath @a path
          insert LoadRec e
          pure Nothing
        Just (Nothing, True, loadRec) -> do
          when (b' && not loadRec) $ insert LoadRec e
          pure Nothing
    Nothing -> do
      e <- loadByRelPath @a path
      when b' $ insert LoadRec e
      pure Nothing
getAsset (SubRef e f) b prev = fmap f <$> getAsset e b prev
getAsset (TransRef a f) b prev = do
  x <- getAsset a b prev
  case x of
    Nothing -> pure Nothing
    Just x -> do
      getAsset (f x) b prev
getAsset (PureRef a) _ _ = pure $ Just a
getAsset (AppRef f a) b prev = do
  x <- getAsset f b prev
  y <- getAsset a b prev
  let z = (,) <$> x <*> y
  case z of
    Nothing -> pure Nothing
    Just (f, a) -> do
      pure $ Just $ f a
getAsset NullRef _ _ = pure Nothing

-- loaded :: forall a. AssetRef a -> System Bool
-- loaded ref = isJust <$> getAsset ref

instance Show (AssetRef a) where
  show (Ref p) = "Ref " ++ show p
  show (EntityRef p) = "Ref " ++ show p
  show (SubRef a _) = "SubRef " ++ show a
  show (TransRef a _) = "TransRef " ++ show a
  show (PureRef _) = "PureRef"
  show (AppRef f a) = "AppRef" ++ show f ++ " " ++ show a
  show NullRef = "NullRef"

instance Functor AssetRef where
  fmap f a = SubRef a f

instance Applicative AssetRef where
  pure :: a -> AssetRef a
  pure = PureRef
  (<*>) :: AssetRef (a -> b) -> AssetRef a -> AssetRef b
  (<*>) = AppRef

instance Monad AssetRef where
  (>>=) :: AssetRef a -> (a -> AssetRef b) -> AssetRef b
  (>>=) = TransRef

class ToAssetRef r where
  toRef :: (Asset a, Bundle a) => r a -> AssetRef a

instance (Asset a) => IsString (AssetRef a) where
  fromString :: String -> AssetRef a
  fromString s = Ref (T.pack s)