{-|
Description: VFS persisted purely through 'StateT'

This is intended primarily to be a drop-in testing mock for VFS implementations, but you may also find it useful within your program if you want to
work with some hierarchically-organized data.

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.Pure
	( PureVFS
	, runPureVFS
	, runPureVFS'
	, runPureVFS_
	) where

import Control.Monad.Trans.State.Lazy
import ClassyPrelude hiding (ByteString, handle, hash)
import Data.Conduit.VFS.Import
import qualified Data.HashMap.Lazy as HashMap
import Data.HashMap.Lazy (HashMap)
import Control.Monad.Extra (whenJust, maybeM)
import Control.Monad.Fail (MonadFail)
import System.FilePath (splitPath, (</>))
import qualified Data.Text as Text

-- | The possible kinds of nodes
data PNode
	= PNodeDir (HashMap Text PNode) -- ^ A directory
	| PNodeFile ByteString          -- ^ A file

-- | The basic implementation of the VFS.
newtype PureVFS m a = PureVFS
	{ unPVFS :: StateT PNode m a }
	deriving (Applicative, Functor, MonadFail, Monad, MonadTrans)

-- | Given a 'PureVFS', run it in the local monad and return both the monadic return value and the root node of the VFS.
runPureVFS :: PureVFS m a -> m (a, PNode)
runPureVFS pvfs = runStateT (unPVFS pvfs) (PNodeDir mempty)
{-# INLINE runPureVFS #-}

-- | Given a 'PureVFS', run it in the local monad and return the root node of the VFS.
runPureVFS' :: (Monad m) => PureVFS m a -> m PNode
runPureVFS' pvfs = execStateT (unPVFS pvfs) (PNodeDir mempty)
{-# INLINE runPureVFS' #-}

-- | Given a 'PureVFS', run it in the local monad and disregard any results.
runPureVFS_ :: (Monad m) => PureVFS m a -> m ()
runPureVFS_ = void . runPureVFS
{-# INLINE runPureVFS_ #-}

-- | Retrieves the root of the 'PureVFS'.
getRoot :: (Monad m) => PureVFS m PNode
getRoot = PureVFS get
{-# INLINE getRoot #-}

-- | Sets the root of the 'PureVFS' and returns the new root
setRoot :: (Monad m) => PNode -> PureVFS m PNode
setRoot newRoot = setRoot_ newRoot >> return newRoot
{-# INLINE setRoot #-}

-- | Sets the root of the 'PureVFS' without returning it
setRoot_ :: (Monad m) => PNode -> PureVFS m ()
setRoot_ = PureVFS . put
{-# INLINE setRoot_ #-}

-- | Changes the root without providing any return value
modifyRoot_ :: (Monad m) => (PNode -> PureVFS m PNode) -> PureVFS m ()
modifyRoot_ f = getRoot >>= f >>= PureVFS . put -- Oh, right: StateT doesn't have concurrency problems.  Sweet!
{-# INLINE modifyRoot_ #-}

-- | Changes the root and provides the updated value.
modifyRoot :: (Monad m) => (PNode -> PureVFS m PNode) -> PureVFS m PNode
modifyRoot f = modifyRoot_ f >> getRoot
{-# INLINE modifyRoot #-}

getNodeC :: (Monad m) => FilePath -> ConduitT i o (PureVFS m) (Maybe PNode)
getNodeC = lift . getNode
{-# INLINE getNodeC #-}

getNode :: (Monad m) => FilePath -> PureVFS m (Maybe PNode)
getNode filepath = loop (splitPath filepath) <$> getRoot
	where
		loop [] _ = Nothing
		loop _ (PNodeFile _) = Nothing
		loop [filename] (PNodeDir hash) = HashMap.lookup (Text.pack filename) hash
		loop (dirname:rest) (PNodeDir hash) = HashMap.lookup (Text.pack dirname) hash >>= loop rest
{-# INLINE getNode #-}

modifyNodeC_ :: (Monad m) => FilePath -> (Maybe PNode -> PureVFS m (Maybe PNode)) -> ConduitT i o (PureVFS m) ()
modifyNodeC_ filepath f = lift $ modifyNode_ filepath f
{-# INLINE modifyNodeC_ #-}

modifyNode_ :: (Monad m) => FilePath -> (Maybe PNode -> PureVFS m (Maybe PNode)) -> PureVFS m ()
modifyNode_ filepath f = modifyRoot_ $ loop (splitPath filepath)
	where
		loop [] node = return node
		loop _ file@(PNodeFile _) = return file
		loop (nodename:rest) dir@(PNodeDir hash) = case HashMap.lookup (Text.pack nodename) hash of
			Nothing -> case rest of
				[] ->
					maybeM
						(return dir)
						(\result -> return . PNodeDir $ HashMap.insert (Text.pack nodename) result hash)
						(f (Just dir))
				(restHead:restTail) ->
					loop restTail (PNodeDir mempty) >>= \result ->
						return . PNodeDir $ HashMap.insert (Text.pack nodename) (PNodeDir $ HashMap.singleton (Text.pack restHead) result) hash
			(Just node) -> loop rest node >>= \result -> return $ PNodeDir $ HashMap.insert (Text.pack nodename) result hash
{-# INLINEABLE modifyNode_ #-}

modifyNodeC :: (Monad m) => FilePath -> (Maybe PNode -> PureVFS m (Maybe PNode)) -> ConduitT i o (PureVFS m) (Maybe PNode)
modifyNodeC filepath f = lift $ modifyNode filepath f
{-# INLINE modifyNodeC #-}

-- | Updates the node at the given filepath and then returns the updated node.
modifyNode :: (Monad m) => FilePath -> (Maybe PNode -> PureVFS m (Maybe PNode)) -> PureVFS m (Maybe PNode)
modifyNode filepath f = modifyNode_ filepath f >> getNode filepath
{-# INLINE modifyNode #-}

instance (Monad m) => ReadVFSC (PureVFS m) where

	vfsTypeC = awaitForever $ \filepath -> do
			maybeNode <- getNodeC filepath
			yield (filepath, toType <$> maybeNode)
		where
			toType (PNodeDir _) = VDirectory
			toType (PNodeFile _) = VFile
	{-# INLINE vfsTypeC #-}

	vfsContentsC = awaitForever $ \filepath -> do
		maybeResult <- getNodeC filepath
		whenJust maybeResult $ \case
			(PNodeFile bytes) -> yield (filepath, bytes)
			(PNodeDir _)      -> return ()
	{-# INLINE vfsContentsC #-}

	vfsChildrenC = awaitForever $ \filepath -> do
		maybeNode <- getNodeC filepath
		case maybeNode of
			Nothing -> yield filepath
			(Just (PNodeFile _)) -> yield filepath
			(Just (PNodeDir hash)) -> yieldMany $ (filepath </>) . Text.unpack <$> HashMap.keys hash
	{-# INLINE vfsChildrenC #-}

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

	vfsWriteSink = awaitForever $ \(filepath, bs) -> modifyNodeC_ filepath (const . return . Just $ PNodeFile bs)
	{-# INLINE vfsWriteSink #-}

	vfsRemoveSink = awaitForever $ flip modifyNodeC_ (const $ return Nothing)
	{-# INLINE vfsRemoveSink #-}

-- | A class denoting that the type is usable as VFS conduits for reading and writing.
instance (Monad m) => VFSC (PureVFS m)