XFESystem

using System.Text;
using DiscUtils;
using DiscUtils.ExFat;
using DiscUtils.Fat;
using DiscUtils.Ntfs;
using XFESystem.Core;

namespace XFESystem.Storage;

public sealed class MountedVolume : IDisposable
{
    const int MaxFileBytes=4*1024*1024;
    readonly IPhysicalDisk disk;
    readonly DiskPartition partition;
    readonly string kind;
    readonly VolumeDirtyState? dirtyState;
    DiscFileSystem? filesystem;
    public string Path {get;}
    public string FileSystem=>kind;
    public string Label {get;}
    public bool Writable {get;private set;}
    public string ReadOnlyReason {get;private set;}="";
    public MountedVolume(IPhysicalDisk disk,DiskPartition partition,string path)
    {
        this.disk=disk;this.partition=partition;Path=SystemPaths.Normalize(path);
        byte[] boot=new byte[512];disk.Read(partition.Start,boot);
        kind=boot.AsSpan(3,8).SequenceEqual("NTFS    "u8)?"NTFS":boot.AsSpan(3,8).SequenceEqual("EXFAT   "u8)?"exFAT":boot.AsSpan(3,8).SequenceEqual("-FVE-FS-"u8)?throw new IOException("BitLocker volume is locked."):"FAT";
        filesystem=Open(false);Label=filesystem.VolumeLabel??"";
        Writable=disk.CanWrite&&partition.DataPartition;
        if(!Writable)ReadOnlyReason="Protected partition or read-only device";
        if(Writable)
        {
            try
            {
                dirtyState=VolumeDirtyState.Create(new PartitionStream(disk,partition,true),kind,filesystem as NtfsFileSystem);
                if(dirtyState.IsDirty){Writable=false;ReadOnlyReason="Filesystem requires recovery in its original operating system";}
                if(filesystem is NtfsFileSystem ntfs)
                {
                    if((ntfs.VolumeInfo.Flags&VolumeInformationFlags.Dirty)!=0)
                    {Writable=false;ReadOnlyReason="NTFS requires recovery in Windows before writing";}
                    else if(Hibernated(ntfs))
                    {Writable=false;ReadOnlyReason="Windows is hibernated; fully shut down Windows before writing";}
                }
            }
            catch(IOException error){Writable=false;ReadOnlyReason=error.Message;}
        }
    }
    DiscFileSystem Open(bool write)
    {
        Stream stream=new PartitionStream(disk,partition,write);
        // The exFAT library otherwise inherits host Path separators. Our guest
        // uses a Linux-ABI BCL while the workspace API always uses backslashes.
        DiscFileSystem fs=kind switch {"NTFS"=>new NtfsFileSystem(stream),"exFAT"=>new ExFatFileSystem(stream,['\\','/']),_=>new FatFileSystem(stream)};
        if(fs is NtfsFileSystem ntfs){ntfs.NtfsOptions.HideHiddenFiles=false;ntfs.NtfsOptions.HideSystemFiles=false;ntfs.NtfsOptions.HideMetafiles=true;}
        return fs;
    }
    static bool Hibernated(NtfsFileSystem fs)
    {
        if(!fs.FileExists("hiberfil.sys"))return false;
        using var file=fs.OpenFile("hiberfil.sys",FileMode.Open,FileAccess.Read);Span<byte> header=stackalloc byte[4];
        if(file.Read(header)!=4)return true;
        return System.Buffers.Binary.BinaryPrimitives.ReadUInt32LittleEndian(header)!=0;
    }
    DiscFileSystem Fs=>filesystem??=Open(false);
    string Relative(string path)
    {
        path=SystemPaths.Normalize(path);
        if(path.Equals(Path,StringComparison.OrdinalIgnoreCase))return "";
        if(!path.StartsWith(Path+"\\",StringComparison.OrdinalIgnoreCase))throw new IOException("Path leaves its mounted volume.");
        string relative=path[(Path.Length+1)..];string current="";
        foreach(string part in relative.Split('\\'))
        {
            if(part.StartsWith('$'))throw new UnauthorizedAccessException("Filesystem metadata is protected.");
            current=current.Length==0?part:current+"\\"+part;
            if(Fs.Exists(current)&&(Fs.GetAttributes(current)&FileAttributes.ReparsePoint)!=0)throw new IOException("Reparse points are not followed by this system.");
        }
        return relative;
    }
    public bool Exists(string path){string relative=Relative(path);return relative.Length==0||Fs.Exists(relative);}
    public WorkspaceEntry[] List(string path)
    {
        string relative=Relative(path);
        return Fs.GetFileSystemEntries(relative).Take(32769).Select(name=>
        {
            string clean=name.TrimStart('\\');string full=Path+"\\"+clean;
            var attributes=Fs.GetAttributes(clean);bool directory=(attributes&FileAttributes.Directory)!=0;
            return new WorkspaceEntry(full,directory,directory?0:Fs.GetFileLength(clean));
        }).OrderBy(e=>!e.IsDirectory).ThenBy(e=>e.Name,StringComparer.OrdinalIgnoreCase).ToArray() is var entries&&entries.Length<=32768?entries:throw new IOException("Directory exceeds the supported entry count.");
    }
    public byte[] ReadBytes(string path)
    {
        string relative=Relative(path);long length=Fs.GetFileLength(relative);
        if(length<0||length>MaxFileBytes)throw new IOException("The current file read limit is 4 MiB.");
        using var file=Fs.OpenFile(relative,FileMode.Open,FileAccess.Read);byte[] bytes=new byte[(int)length];file.ReadExactly(bytes);return bytes;
    }
    public string ReadText(string path)
    {using var reader=new StreamReader(new MemoryStream(ReadBytes(path),false),new UTF8Encoding(false,true),true);return reader.ReadToEnd();}
    void Mutate(Action<DiscFileSystem> action)
    {
        if(!Writable||dirtyState==null)throw new UnauthorizedAccessException("Read-only volume: "+ReadOnlyReason);
        filesystem?.Dispose();filesystem=null;
        try
        {
            dirtyState.Set(true);
            using(var fs=Open(true))action(fs);
            disk.Flush();dirtyState.Set(false);
        }
        catch {Writable=false;ReadOnlyReason="A filesystem operation failed; the dirty flag was retained for recovery";throw;}
    }
    public void WriteBytes(string path,ReadOnlySpan<byte> content)
    {
        string relative=Relative(path);if(relative.Length==0||content.Length>MaxFileBytes)throw new IOException("Invalid file path or size.");
        int separator=relative.LastIndexOf('\\');string parent=separator<0?"":relative[..separator];
        if(!Fs.DirectoryExists(parent))throw new DirectoryNotFoundException(parent);
        byte[] bytes=content.ToArray();Mutate(fs=>{using var file=fs.OpenFile(relative,FileMode.Create,FileAccess.Write);file.Write(bytes);file.Flush();});
    }
    public void WriteText(string path,string text)=>WriteBytes(path,Encoding.UTF8.GetBytes(text));
    public void CreateDirectory(string path)
    {string relative=Relative(path);if(relative.Length==0||Fs.DirectoryExists(relative))return;Mutate(fs=>fs.CreateDirectory(relative));}
    public void Delete(string path)
    {
        string relative=Relative(path);if(relative.Length==0)throw new IOException("A volume root cannot be deleted.");
        bool directory=Fs.DirectoryExists(relative);if(!directory&&!Fs.FileExists(relative))throw new FileNotFoundException(path);
        Mutate(fs=>{if(directory)fs.DeleteDirectory(relative,true);else fs.DeleteFile(relative);});
    }
    public void Move(string source,string destination)
    {
        string from=Relative(source),to=Relative(destination);
        if(from.Length==0||to.Length==0)throw new IOException("A volume root cannot be moved.");
        bool directory=Fs.DirectoryExists(from);if(!directory&&!Fs.FileExists(from))throw new FileNotFoundException(source);
        if(Fs.Exists(to))throw new IOException("Destination already exists.");
        Mutate(fs=>{if(directory)fs.MoveDirectory(from,to);else fs.MoveFile(from,to);});
    }
    public void Dispose(){filesystem?.Dispose();filesystem=null;}
}