{-|
Description: VFS persisted in-memory

This is intended to demonstrate how a VFS can be defined, but is also entirely usable in code if you find its functionality useful.

This VFS implementation stores its values in a lazy in-memory map, which is itself persisted into an 'MVar' and captured in the monadic state.
This requires the monad to be in the class 'MonadUnliftIO', but that allows the state to be shared throughout the application, maintaining
value consistency across threads and even across invocations.

Individual read and write operations are atomic.

The 'FilePath' values used in this VFS are split using 'splitPath' and joined using '</>', but are otherwise used directly: there is no concept of
paths being "relative" or "absolute" for this VFS.  It is also possible for a file and a directory to have the same name, since directories names
are appended with @/@, as per 'splitPath'.  (This implementation detail is up for debate and may be changed in a future major release: please file
an issue if you want to have a discussion around it.)

-}
module Data.Conduit.VFS.InMemory
        ( InMemoryVFS
        , InMemoryVFSRoot
        , runInMemoryVFS
        , runInMemoryVFS'
        , mkInMemoryVFSRoot
        ) where

import ClassyPrelude hiding (ByteString, handle)
import Control.Monad.Extra (ifM)
import Data.Conduit.VFS.Import
import qualified Data.ByteString.Lazy as LBS
import qualified Data.HashMap.Lazy as HashMap
import Data.HashMap.Lazy (HashMap)
import System.FilePath (splitPath)
import Control.Monad.Fail (MonadFail)
import qualified Data.Text as Text

-- | The possible kinds of nodes
data IMNode
        = IMNodeDir IMDirectory -- ^ VDirectory
        | IMNodeFile IMFile     -- ^ VFile

-- | How is the data stored?
data IMFile
        = Resident ByteString  -- ^ Data is stored resident in memory
        | EmptyFile            -- ^ Data is empty

-- | Definition of a directory
newtype IMDirectory = IMDirectory
        { imdNodes :: HashMap Text IMNode   -- ^ The nodes contained within the directory
        }

instance Semigroup IMDirectory where
        -- | Right-biased, but with directories recursively merged.
        (<>) imdLeft imdRight =
                        IMDirectory { imdNodes = HashMap.unionWith mergeImpl nodeLeft nodeRight }
                where
                        nodeLeft = imdNodes imdLeft
                        nodeRight = imdNodes imdRight
                        mergeImpl (IMNodeDir nodeDirLeft) (IMNodeDir nodeDirRight) = IMNodeDir $ nodeDirLeft <> nodeDirRight -- Recursively merge two dirs
                        mergeImpl _ right = right -- Default to the right if there's any other kind of conflict.
        {-# INLINE (<>) #-}

instance Monoid IMDirectory where
        mempty = IMDirectory { imdNodes = mempty }
        {-# INLINE mempty #-}

type instance Element IMDirectory = (Text, IMNode)

instance MonoFunctor IMDirectory where
        omap f oldImd = IMDirectory{ imdNodes = HashMap.fromList (f <$> HashMap.toList (imdNodes oldImd)) }
        {-# INLINE omap #-}

{- TODO Implement the MonoFoldable methods.
instance MonoFoldable IMDirectory where
	otoList IMDirectory{imdNodes} = toList imdNodes
	{-# INLINE otoList #-}
-}

instance MonoPointed IMDirectory where
        opoint (k,v) = IMDirectory { imdNodes = HashMap.singleton k v }
        {-# INLINE opoint #-}

-- | The root of the VFS.
newtype InMemoryVFSRoot = InMemoryVFSRoot { imvfsStore :: MVar IMDirectory }

-- | The basic implementation of the VFS.
newtype InMemoryVFS m a = InMemoryVFS
        { unIMVFS :: ReaderT InMemoryVFSRoot m a }
        deriving (Functor, Applicative, Monad, MonadTrans, MonadIO, MonadFail, MonadReader InMemoryVFSRoot)

-- | Creates an 'InMemoryVFSRoot' that can be shared among many 'InMemoryVFS' invocations.
mkInMemoryVFSRoot :: (MonadIO m) => m InMemoryVFSRoot
mkInMemoryVFSRoot = do
        mvar <- newMVar mempty
        return $ InMemoryVFSRoot { imvfsStore = mvar }
{-# INLINE mkInMemoryVFSRoot #-}

-- | Given an 'InMemoryVFS', run it in the local monad.
runInMemoryVFS :: (MonadUnliftIO m) => InMemoryVFS m a -> m a
runInMemoryVFS imvfs = mkInMemoryVFSRoot >>= flip runInMemoryVFS' imvfs
{-# INLINE runInMemoryVFS #-}

-- | Runs an 'InMemoryVFS' using a provided 'InMemoryVFSRoot'.
runInMemoryVFS' :: InMemoryVFSRoot -> InMemoryVFS m a -> m a
runInMemoryVFS' root imvfs =
        let monad = unIMVFS imvfs in
        runReaderT monad root
{-# INLINE runInMemoryVFS' #-}

-- | Takes a function that consumes the root dir and produces a monadic action, and then returns that monadic
--   action as an 'InMemoryVFS'.
withIMVFSRootDir :: (MonadUnliftIO m) => ( IMDirectory -> m a ) -> InMemoryVFS m a
withIMVFSRootDir f = do
                mvar <- imvfsStore <$> ask
                rootDir <- readMVar mvar
                lift $ f rootDir
{-# INLINE withIMVFSRootDir #-}

-- | Takes a function that consumes the root dir, produces a new root dir as a monadic action, and then returns
--   that monadic action as an 'InMemoryVFS'.
modifyIMVFSRootDir :: (MonadUnliftIO m) => ( IMDirectory -> m IMDirectory ) -> InMemoryVFS m ()
modifyIMVFSRootDir f = do
        mvar <- imvfsStore <$> ask
        lift $ modifyMVar_ mvar f
{-# INLINE modifyIMVFSRootDir #-}

instance (MonadUnliftIO m) => ReadVFSC (InMemoryVFS m) where

        vfsTypeC = awaitForever $ \filepath -> do
                        result <- lift $ withIMVFSRootDir $ return . loop (splitPath filepath)
                        yield (filepath, result)
                where
                        loop :: [FilePath] -> IMDirectory -> Maybe VFileType
                        loop [] (IMDirectory _) = Just VDirectory
                        loop ("/":rest) imd = loop rest imd
                        loop (nextDirPath:rest) IMDirectory{imdNodes} =
                                case HashMap.lookup (Text.pack nextDirPath) imdNodes of
                                        Nothing           -> Nothing
                                        (Just imnode)     ->
                                                case imnode of
                                                        (IMNodeDir imd) -> loop rest imd
                                                        (IMNodeFile _)
                                                                | null rest   -> Just VFile
                                                                | otherwise   -> Nothing
        {-# INLINEABLE vfsTypeC #-}

        vfsContentsC = awaitForever $ \filepath -> do
                        result <- lift $ withIMVFSRootDir $ loop (splitPath filepath)
                        case result of
                                Nothing          -> return ()
                                (Just resultLBS) -> yield (filepath, resultLBS)
                where
                        loop :: [FilePath] -> IMDirectory -> m (Maybe ByteString)
                        loop [] _ = return Nothing
                        loop [filename] IMDirectory{imdNodes} =
                                case HashMap.lookup (Text.pack filename) imdNodes of
                                        Nothing                     -> return Nothing
                                        (Just (IMNodeDir _))        -> return Nothing
                                        (Just (IMNodeFile fileData)) ->
                                                case fileData of
                                                        (Resident bytes)    -> return (Just bytes)
                                                        EmptyFile           -> return (Just mempty)
                        loop (dirname:rest) IMDirectory{imdNodes} =
                                case HashMap.lookup (Text.pack dirname) imdNodes of
                                        (Just (IMNodeDir imd)) -> loop rest imd
                                        _                      -> return Nothing
        {-# INLINEABLE vfsContentsC #-}

        vfsContentsEitherC = awaitForever $ \filepath -> do
                        maybeResult <- lift $ withIMVFSRootDir $ loop (splitPath filepath)
                        case maybeResult of
                                Nothing -> return ()
                                (Just bytes) -> do
                                        yield $ Left filepath
                                        yield $ Right bytes
                where
                        loop :: [FilePath] -> IMDirectory -> m (Maybe ByteString)
                        loop [] _ = return Nothing
                        loop [filename] IMDirectory{imdNodes} =
                                case HashMap.lookup (Text.pack filename) imdNodes of
                                        Nothing                      -> return Nothing
                                        (Just (IMNodeDir _))         -> return Nothing
                                        (Just (IMNodeFile filedata)) ->
                                                case filedata of
                                                        (Resident bytes)    -> return $ Just bytes
                                                        EmptyFile           -> return $ Just mempty
                        loop (dirname:rest) IMDirectory{imdNodes} =
                                case HashMap.lookup (Text.pack dirname) imdNodes of
                                        (Just (IMNodeDir imd)) -> loop rest imd
                                        _                      -> return Nothing
        {-# INLINEABLE vfsContentsEitherC #-}

        vfsChildrenC = awaitForever $ \filepath ->
                        lift ( withIMVFSRootDir $ return . loop filepath (splitPath filepath) ) >>= yieldMany
                where
                        loop :: FilePath -> [FilePath] -> IMDirectory -> [FilePath]
                        loop _ [] IMDirectory{imdNodes} = Text.unpack <$> HashMap.keys imdNodes
                        loop filepath (foo:rest) IMDirectory{imdNodes} =
                                case HashMap.lookup (Text.pack foo) imdNodes of
                                        (Just (IMNodeDir imd@(IMDirectory dir)))
                                                | null rest          -> (filepath </>) . Text.unpack <$> HashMap.keys dir
                                                | otherwise          -> loop filepath rest imd
                                        _                      -> mempty
        {-# INLINEABLE vfsChildrenC #-}

-- | A class denoting that the type is usable as VFS conduits for writing.
instance (MonadUnliftIO m) => WriteVFSC (InMemoryVFS m) where

        vfsWriteEitherSink = awaitForever $ \case
                (Right _) -> return () -- Ignore: bytes without a file they belong to!
                (Left filepath) -> awaitBytes >>= \bytes ->
                        let imfile =
                                if null bytes then
                                        EmptyFile
                                else
                                        Resident bytes
                        in
                        lift $ modifyIMVFSRootDir $ return . loop (IMNodeFile imfile) (splitPath filepath)
                where
                        hasMoreBytes = peekC >>= \case
                                (Just (Right _)) -> return True
                                _ -> return False
                        awaitBytes =
                                flip (ifM hasMoreBytes) (return mempty) $
                                        await >>= \case
                                                (Just (Right bytes)) -> LBS.append bytes <$> awaitBytes
                                                _ -> fail "We should have more bytes, but we don't."
                        loop _ [] imd = imd
                        loop node [filename] imd@IMDirectory{imdNodes} = imd { imdNodes = HashMap.insert (Text.pack filename) node imdNodes }
                        loop node (name:rest) imd@IMDirectory{imdNodes} = imd { imdNodes = HashMap.alter
                                (\case
                                        Nothing                     -> Just . IMNodeDir . loop node rest $ IMDirectory { imdNodes=mempty }
                                        (Just (IMNodeDir childImd)) -> Just . IMNodeDir $ loop node rest childImd
                                        whatever                    -> whatever
                                ) (Text.pack name) imdNodes }
        {-# INLINEABLE vfsWriteEitherSink #-}

        vfsRemoveSink = awaitForever $ \filepath -> lift . modifyIMVFSRootDir $ return . loop (splitPath filepath)
                where
                        loop [] imd = imd
                        loop [filename] imd@IMDirectory{imdNodes} = imd { imdNodes = HashMap.delete (Text.pack filename) imdNodes }
                        loop (name:rest) imd@IMDirectory{imdNodes} = imd
                                { imdNodes = HashMap.adjust
                                        (\case
                                                (IMNodeDir childImd) -> IMNodeDir $ loop rest childImd
                                                whatever             -> whatever
                                        ) (Text.pack name) imdNodes
                                }
        {-# INLINE vfsRemoveSink #-}

instance (MonadUnliftIO m) => VFSC (InMemoryVFS m)