InMemoryFileSystem.cs

// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System.Text;
using System.Text.RegularExpressions;
using Microsoft.TemplateEngine.Abstractions.PhysicalFileSystem;

namespace Microsoft.TemplateEngine.Utils
{
    /// <summary>
    /// In-memory file system implementation of <see cref="IPhysicalFileSystem"/>.
    /// </summary>
    /// <seealso cref="Abstractions.ITemplateEngineHost"/>
    public class InMemoryFileSystem : IPhysicalFileSystem
    {
        private readonly FileSystemDirectory _root;
        private readonly IPhysicalFileSystem _basis;

        public InMemoryFileSystem(string root, IPhysicalFileSystem basis)
        {
            _basis = basis;
            _root = new FileSystemDirectory(Path.GetFileName(root.TrimEnd('/', '\\')), root);
            _ = IsPathInCone(root, out string newRoot);
            if (root != newRoot)
            {
                _root = new FileSystemDirectory(Path.GetFileName(newRoot.TrimEnd('/', '\\')), newRoot);
            }
        }

        public void CreateDirectory(string path)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                _basis.CreateDirectory(path);
                return;
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');

            if (string.IsNullOrEmpty(rel))
            {
                return;
            }

            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }
        }

        public Stream CreateFile(string path)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                return _basis.CreateFile(path);
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            if (!currentDir.Files.TryGetValue(parts[parts.Length - 1], out FileSystemFile targetFile))
            {
                targetFile = new FileSystemFile(parts[parts.Length - 1], Path.Combine(currentDir.FullPath, parts[parts.Length - 1]));
                currentDir.Files[parts[parts.Length - 1]] = targetFile;
            }

            return targetFile.OpenWrite();
        }

        public void DirectoryDelete(string path, bool recursive)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                //TODO: handle cases where a portion of what would be deleted is inside our cone
                //  and parts are outside
                _basis.DirectoryDelete(path, recursive);
                return;
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');

            if (string.IsNullOrEmpty(rel))
            {
                return;
            }

            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;
            FileSystemDirectory? parent = null;

            for (int i = 0; i < parts.Length; ++i)
            {
                parent = currentDir;
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    return;
                }

                currentDir = dir;
            }

            if (!recursive && (currentDir.Directories.Count > 0 || currentDir.Files.Count > 0))
            {
                throw new IOException("Directory is not empty");
            }

            _ = (parent?.Directories.Remove(currentDir.Name));
        }

        public bool DirectoryExists(string directory)
        {
            if (!IsPathInCone(directory, out string processedPath))
            {
                return _basis.DirectoryExists(directory);
            }

            directory = processedPath;
            string rel = directory.Substring(_root.FullPath.Length).Trim('/', '\\');

            if (string.IsNullOrEmpty(rel))
            {
                return true;
            }

            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    return false;
                }

                currentDir = dir;
            }

            return true;
        }

        public IEnumerable<string> EnumerateDirectories(string path, string pattern, SearchOption searchOption)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                //TODO: Handle cases where part of the directory set is inside our cone and part is not
                foreach (string s in _basis.EnumerateDirectories(path, pattern, searchOption))
                {
                    yield return s;
                }

                yield break;
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');
            FileSystemDirectory currentDir = _root;

            if (!string.IsNullOrEmpty(rel))
            {
                string[] parts = rel.Split('/', '\\');
                for (int i = 0; i < parts.Length; ++i)
                {
                    if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                    {
                        yield break;
                    }

                    currentDir = dir;
                }
            }

            Regex rx = new(Regex.Escape(pattern).Replace("\\*", ".*").Replace("\\?", "."));

            if (searchOption == SearchOption.TopDirectoryOnly)
            {
                foreach (KeyValuePair<string, FileSystemDirectory> entry in currentDir.Directories)
                {
                    if (rx.IsMatch(entry.Key))
                    {
                        yield return entry.Value.FullPath;
                    }
                }

                yield break;
            }

            Stack<IEnumerator<KeyValuePair<string, FileSystemDirectory>>> directories = new();
            IEnumerator<KeyValuePair<string, FileSystemDirectory>> current = currentDir.Directories.GetEnumerator();
            bool moveNext;

            while ((moveNext = current.MoveNext()) || directories.Count > 0)
            {
                while (!moveNext)
                {
                    current.Dispose();

                    if (directories.Count == 0)
                    {
                        break;
                    }

                    current = directories.Pop();
                    moveNext = current.MoveNext();
                }

                if (!moveNext)
                {
                    break;
                }

                if (rx.IsMatch(current.Current.Key))
                {
                    yield return current.Current.Value.FullPath;
                }

                directories.Push(current);
                current = current.Current.Value.Directories.GetEnumerator();
            }
        }

        public IEnumerable<string> EnumerateFiles(string path, string pattern, SearchOption searchOption)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                //TODO: Handle cases where part of the directory set is inside our cone and part is not
                foreach (string s in _basis.EnumerateFiles(path, pattern, searchOption))
                {
                    yield return s;
                }

                yield break;
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');
            FileSystemDirectory currentDir = _root;

            if (!string.IsNullOrEmpty(rel))
            {
                string[] parts = rel.Split('/', '\\');
                for (int i = 0; i < parts.Length; ++i)
                {
                    if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                    {
                        yield break;
                    }

                    currentDir = dir;
                }
            }

            Regex rx = new(Regex.Escape(pattern).Replace("\\*", ".*").Replace("\\?", "."));
            foreach (KeyValuePair<string, FileSystemFile> entry in currentDir.Files)
            {
                if (rx.IsMatch(entry.Key))
                {
                    yield return entry.Value.FullPath;
                }
            }

            if (searchOption == SearchOption.TopDirectoryOnly)
            {
                yield break;
            }

            Stack<IEnumerator<KeyValuePair<string, FileSystemDirectory>>> directories = new();
            IEnumerator<KeyValuePair<string, FileSystemDirectory>> current = currentDir.Directories.GetEnumerator();
            bool moveNext;

            while ((moveNext = current.MoveNext()) || directories.Count > 0)
            {
                while (!moveNext)
                {
                    current.Dispose();

                    if (directories.Count == 0)
                    {
                        break;
                    }

                    current = directories.Pop();
                    moveNext = current.MoveNext();
                }

                if (!moveNext)
                {
                    break;
                }

                foreach (KeyValuePair<string, FileSystemFile> entry in current.Current.Value.Files)
                {
                    if (rx.IsMatch(entry.Key))
                    {
                        yield return entry.Value.FullPath;
                    }
                }

                directories.Push(current);
                current = current.Current.Value.Directories.GetEnumerator();
            }
        }

        public IEnumerable<string> EnumerateFileSystemEntries(string directoryName, string pattern, SearchOption searchOption)
        {
            if (!IsPathInCone(directoryName, out string processedPath))
            {
                //TODO: Handle cases where part of the directory set is inside our cone and part is not
                foreach (string s in _basis.EnumerateFileSystemEntries(directoryName, pattern, searchOption))
                {
                    yield return s;
                }

                yield break;
            }

            directoryName = processedPath;
            string rel = directoryName.Substring(_root.FullPath.Length).Trim('/', '\\');
            FileSystemDirectory currentDir = _root;

            if (!string.IsNullOrEmpty(rel))
            {
                string[] parts = rel.Split('/', '\\');
                for (int i = 0; i < parts.Length; ++i)
                {
                    if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                    {
                        yield break;
                    }

                    currentDir = dir;
                }
            }

            Regex rx = new("^" + Regex.Escape(pattern).Replace("\\*", ".*").Replace("\\?", ".") + "$");

            foreach (KeyValuePair<string, FileSystemFile> entry in currentDir.Files)
            {
                if (rx.IsMatch(entry.Key))
                {
                    yield return entry.Value.FullPath;
                }
            }

            if (searchOption == SearchOption.TopDirectoryOnly)
            {
                foreach (KeyValuePair<string, FileSystemDirectory> entry in currentDir.Directories)
                {
                    if (rx.IsMatch(entry.Key))
                    {
                        yield return entry.Value.FullPath;
                    }
                }
                yield break;
            }

            Stack<IEnumerator<KeyValuePair<string, FileSystemDirectory>>> directories = new();
            IEnumerator<KeyValuePair<string, FileSystemDirectory>> current = currentDir.Directories.GetEnumerator();
            bool moveNext;

            while ((moveNext = current.MoveNext()) || directories.Count > 0)
            {
                while (!moveNext)
                {
                    current.Dispose();

                    if (directories.Count == 0)
                    {
                        break;
                    }

                    current = directories.Pop();
                    moveNext = current.MoveNext();
                }

                if (!moveNext)
                {
                    break;
                }

                if (rx.IsMatch(current.Current.Key))
                {
                    yield return current.Current.Value.FullPath;
                }

                foreach (KeyValuePair<string, FileSystemFile> entry in current.Current.Value.Files)
                {
                    if (rx.IsMatch(entry.Key))
                    {
                        yield return entry.Value.FullPath;
                    }
                }

                directories.Push(current);
                current = current.Current.Value.Directories.GetEnumerator();
            }
        }

        public void FileCopy(string sourcePath, string targetPath, bool overwrite)
        {
            if (!overwrite && FileExists(targetPath))
            {
                throw new IOException($"File already exists {targetPath}");
            }

            using Stream s = OpenRead(sourcePath);
            using Stream t = CreateFile(targetPath);
            s.CopyTo(t);
        }

        public void FileDelete(string path)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                _basis.FileDelete(path);
                return;
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            _ = currentDir.Files.Remove(parts[parts.Length - 1]);
        }

        public bool FileExists(string file)
        {
            if (!IsPathInCone(file, out string processedPath))
            {
                return _basis.FileExists(file);
            }

            file = processedPath;
            string rel = file.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    return false;
                }

                currentDir = dir;
            }

            return currentDir.Files.ContainsKey(parts[parts.Length - 1]);
        }

        public string GetCurrentDirectory()
        {
            return Directory.GetCurrentDirectory();
        }

        public Stream OpenRead(string path)
        {
            if (!IsPathInCone(path, out string processedPath))
            {
                return _basis.OpenRead(path);
            }

            path = processedPath;
            string rel = path.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            if (!currentDir.Files.TryGetValue(parts[parts.Length - 1], out FileSystemFile targetFile))
            {
                throw new FileNotFoundException("File not found", path);
            }

            return targetFile.OpenRead();
        }

        public string ReadAllText(string path)
        {
            using Stream s = OpenRead(path);
            using StreamReader r = new(s, Encoding.UTF8, true, 8192, true);
            return r.ReadToEnd();
        }

        public byte[] ReadAllBytes(string path)
        {
            //file can be either memory or physical file here
            using Stream s = OpenRead(path);
            if (s is not MemoryStream ms)
            {
                using MemoryStream stream = new();
                s.CopyTo(stream);
                return stream.ToArray();
            }
            return ms.ToArray();
        }

        public void WriteAllText(string path, string value)
        {
            using Stream s = CreateFile(path);
            using StreamWriter r = new(s, Encoding.UTF8, 8192, true);
            r.Write(value);
        }

        public FileAttributes GetFileAttributes(string file)
        {
            if (!IsPathInCone(file, out string processedPath))
            {
                return _basis.GetFileAttributes(file);
            }

            file = processedPath;
            string rel = file.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            if (!currentDir.Files.TryGetValue(parts[parts.Length - 1], out FileSystemFile targetFile))
            {
                throw new FileNotFoundException("File not found", file);
            }

            return targetFile.Attributes;
        }

        public void SetFileAttributes(string file, FileAttributes attributes)
        {
            if (!IsPathInCone(file, out string processedPath))
            {
                _basis.SetFileAttributes(file, attributes);
                return;
            }

            file = processedPath;
            string rel = file.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            if (!currentDir.Files.TryGetValue(parts[parts.Length - 1], out FileSystemFile targetFile))
            {
                throw new FileNotFoundException("File not found", file);
            }

            targetFile.Attributes = attributes;
        }

        public DateTime GetLastWriteTimeUtc(string file)
        {
            if (!IsPathInCone(file, out string processedPath))
            {
                return _basis.GetLastWriteTimeUtc(file);
            }

            file = processedPath;
            string rel = file.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            if (!currentDir.Files.TryGetValue(parts[parts.Length - 1], out FileSystemFile targetFile))
            {
                throw new FileNotFoundException("File not found", file);
            }

            return targetFile.LastWriteTimeUtc;
        }

        public void SetLastWriteTimeUtc(string file, DateTime lastWriteTimeUtc)
        {
            if (!IsPathInCone(file, out string processedPath))
            {
                _basis.SetLastWriteTimeUtc(file, lastWriteTimeUtc);
            }

            file = processedPath;
            string rel = file.Substring(_root.FullPath.Length).Trim('/', '\\');
            string[] parts = rel.Split('/', '\\');
            FileSystemDirectory currentDir = _root;

            for (int i = 0; i < parts.Length - 1; ++i)
            {
                if (!currentDir.Directories.TryGetValue(parts[i], out FileSystemDirectory dir))
                {
                    dir = new FileSystemDirectory(parts[i], Path.Combine(currentDir.FullPath, parts[i]));
                    currentDir.Directories[parts[i]] = dir;
                }

                currentDir = dir;
            }

            if (!currentDir.Files.TryGetValue(parts[parts.Length - 1], out FileSystemFile targetFile))
            {
                throw new FileNotFoundException("File not found", file);
            }

            targetFile.LastWriteTimeUtc = lastWriteTimeUtc;
        }

        /// <inheritdoc/>
        public string PathRelativeTo(string target, string relativeTo)
        {
            if (
                !IsPathInCone(target, out string targetProcessed)
                ||
                !IsPathInCone(relativeTo, out string relativeToProcessed))
            {
                return _basis.PathRelativeTo(target, relativeTo);
            }

            return PhysicalFileSystem.PathRelativeToInternal(targetProcessed, relativeToProcessed);
        }

        /// <summary>
        /// Currently not implemented in <see cref="InMemoryFileSystem"/>.
        /// Just returns <see cref="IDisposable"/> object, but never calls callback.
        /// </summary>
        public IDisposable WatchFileChanges(string filePath, FileSystemEventHandler fileChanged)
        {
            return new MemoryStream(); //Just some disposable dummy
        }

        private bool IsPathInCone(string path, out string processedPath)
        {
            if (!Path.IsPathRooted(path))
            {
                path = Path.Combine(GetCurrentDirectory(), path);
            }

            path = path.Replace('\\', '/');

            bool leadSlash = path[0] == '/';

            if (leadSlash)
            {
                path = path.Substring(1);
            }

            string[] parts = path.Split('/');

            List<string> realParts = new();

            for (int i = 0; i < parts.Length; ++i)
            {
                if (string.IsNullOrEmpty(parts[i]))
                {
                    continue;
                }

                switch (parts[i])
                {
                    case ".":
                        continue;
                    case "..":
                        realParts.RemoveAt(realParts.Count - 1);
                        break;
                    default:
                        realParts.Add(parts[i]);
                        break;
                }
            }

            if (leadSlash)
            {
                realParts.Insert(0, string.Empty);
            }

            processedPath = string.Join(Path.DirectorySeparatorChar + string.Empty, realParts);
            if (processedPath.Equals(_root.FullPath) || processedPath.StartsWith(_root.FullPath.TrimEnd('/', '\\') + Path.DirectorySeparatorChar))
            {
                return true;
            }
            else
            {
                return false;
            }
        }

        private class FileSystemDirectory
        {
            public FileSystemDirectory(string name, string fullPath)
            {
                Name = name;
                FullPath = fullPath;
                Directories = new Dictionary<string, FileSystemDirectory>();
                Files = new Dictionary<string, FileSystemFile>();
            }

            public string Name { get; }

            public string FullPath { get; }

            public Dictionary<string, FileSystemDirectory> Directories { get; }

            public Dictionary<string, FileSystemFile> Files { get; }
        }

        private class FileSystemFile
        {
            private readonly ReaderWriterLockSlim _lock = new();
            private byte[] _data;
            private int _currentReaders;
            private int _currentWriters;

            public FileSystemFile(string name, string fullPath)
            {
                Name = name;
                FullPath = fullPath;
                _data = [];
            }

            public string Name { get; }

            public string FullPath { get; }

            public FileAttributes Attributes { get; set; }

            public DateTime LastWriteTimeUtc { get; set; }

            public Stream OpenRead()
            {
                if (_currentWriters > 0)
                {
                    throw new IOException("File is currently locked for writing");
                }

                _lock.EnterReadLock();
                try
                {
                    if (_currentWriters > 0)
                    {
                        throw new IOException("File is currently locked for writing");
                    }

                    ++_currentReaders;
                    return new DisposingStream(new MemoryStream(_data, false), () =>
                            {
                                _lock.EnterWriteLock();
                                try
                                {
                                    --_currentReaders;
                                }
                                finally
                                {
                                    _lock.ExitWriteLock();
                                }
                            });
                }
                finally
                {
                    _lock.ExitReadLock();
                }
            }

            public Stream OpenWrite()
            {
                if (_currentReaders > 0)
                {
                    throw new IOException("File is currently locked for reading");
                }

                if (_currentWriters > 0)
                {
                    throw new IOException("File is currently locked for writing");
                }

                _lock.EnterWriteLock();
                try
                {
                    if (_currentReaders > 0)
                    {
                        throw new IOException("File is currently locked for reading");
                    }

                    if (_currentWriters > 0)
                    {
                        throw new IOException("File is currently locked for writing");
                    }

                    ++_currentWriters;
                    MemoryStream target = new();
                    return new DisposingStream(target, () =>
                    {
                        _lock.EnterWriteLock();
                        try
                        {
                            --_currentWriters;
                            _data = new byte[target.Length];
                            target.Position = 0;
                            _ = target.Read(_data, 0, _data.Length);
                            LastWriteTimeUtc = DateTime.UtcNow;
                        }
                        finally
                        {
                            _lock.ExitWriteLock();
                        }
                    });
                }
                finally
                {
                    _lock.ExitWriteLock();
                }
            }
        }

        private class DisposingStream : Stream
        {
            private readonly Stream _basis;
            private readonly Action _onDispose;
            private bool _isDisposed;

            public DisposingStream(Stream basis, Action onDispose)
            {
                _onDispose = onDispose;
                _basis = basis;
            }

            public override bool CanRead => _basis.CanRead;

            public override bool CanSeek => _basis.CanSeek;

            public override bool CanWrite => _basis.CanWrite;

            public override long Length => _basis.Length;

            public override long Position { get => _basis.Position; set => _basis.Position = value; }

            public override void Flush() => _basis.Flush();

            public override int Read(byte[] buffer, int offset, int count) => _basis.Read(buffer, offset, count);

            public override long Seek(long offset, SeekOrigin origin) => _basis.Seek(offset, origin);

            public override void SetLength(long value) => _basis.SetLength(value);

            public override void Write(byte[] buffer, int offset, int count) => _basis.Write(buffer, offset, count);

            protected override void Dispose(bool disposing)
            {
                if (!_isDisposed)
                {
                    _isDisposed = true;
                    _onDispose();
                }

                base.Dispose(disposing);
            }
        }
    }
}