XFESystem

using System.Buffers.Binary;
using DiscUtils.Ntfs;

namespace XFESystem.Storage;

// DiscUtils performs format-aware metadata operations. Bracket each mutation
// with the filesystem's dirty flag and device-cache barriers. This is not an
// implementation of Windows $LogFile replay or an atomic power-loss journal.
internal sealed class VolumeDirtyState
{
    readonly PartitionStream stream;
    readonly string kind;
    readonly long[] offsets;
    readonly int recordBytes, sectorBytes;
    public bool IsDirty { get; }
    VolumeDirtyState(PartitionStream stream,string kind,long[] offsets,int recordBytes,int sectorBytes,bool dirty)
    {this.stream=stream;this.kind=kind;this.offsets=offsets;this.recordBytes=recordBytes;this.sectorBytes=sectorBytes;IsDirty=dirty;}
    static ushort U16(byte[] b,int at)=>BinaryPrimitives.ReadUInt16LittleEndian(b.AsSpan(at));
    static uint U32(byte[] b,int at)=>BinaryPrimitives.ReadUInt32LittleEndian(b.AsSpan(at));
    static ulong U64(byte[] b,int at)=>BinaryPrimitives.ReadUInt64LittleEndian(b.AsSpan(at));
    static void W16(byte[] b,int at,ushort value)=>BinaryPrimitives.WriteUInt16LittleEndian(b.AsSpan(at),value);
    static byte[] Read(Stream stream,long offset,int count){stream.Position=offset;byte[] b=new byte[count];stream.ReadExactly(b);return b;}
    public static VolumeDirtyState Create(PartitionStream stream,string kind,NtfsFileSystem? ntfs)
    {
        byte[] boot=Read(stream,0,512);
        if(kind=="NTFS")
        {
            int sector=U16(boot,11),cluster=checked(sector*boot[13]);int encoded=(sbyte)boot[64];
            int record=encoded<0&&encoded>=-16?1<<-encoded:encoded>0?checked(encoded*cluster):0;
            if(sector is not(512 or 4096)||record<sector||record>65536||record%sector!=0)throw new IOException("Unsupported NTFS record geometry.");
            var first=ntfs!.PathToExtents("$MFT").First();
            if(first.Start<0||first.Length<4L*record||first.Start>stream.Length-4L*record)throw new IOException("The initial NTFS MFT records are not in a supported extent.");
            long mirror=checked((long)U64(boot,56)*cluster);
            if(mirror<0||mirror>stream.Length-4L*record)throw new IOException("NTFS mirror leaves its partition.");
            long[] positions=[first.Start+3L*record,mirror+3L*record];bool dirty=false;
            foreach(long at in positions){byte[] raw=Read(stream,at,record);int flags=NtfsFlags(raw,sector);dirty|=(U16(raw,flags)&1)!=0;}
            return new(stream,kind,positions,record,sector,dirty);
        }
        if(kind=="exFAT")
        {
            int sector=1<<boot[108];if(sector is not(512 or 4096))throw new IOException("Unsupported exFAT sector size.");
            byte[] backup=Read(stream,12L*sector,512);if(!backup.AsSpan(3,8).SequenceEqual("EXFAT   "u8))throw new IOException("Invalid exFAT backup boot sector.");
            return new(stream,kind,[0,12L*sector],512,sector,((U16(boot,106)|U16(backup,106))&2)!=0);
        }
        int bytesPerSector=U16(boot,11),reserved=U16(boot,14),fats=boot[16];
        uint fatSectors=U16(boot,22);bool fat32=fatSectors==0;if(fat32)fatSectors=U32(boot,36);
        if(bytesPerSector is not(512 or 4096)||reserved==0||fats<1||fats>2||fatSectors==0)throw new IOException("Unsupported FAT geometry.");
        if(!fat32)
        {
            uint total=U16(boot,19);if(total==0)total=U32(boot,32);
            uint roots=(uint)((U16(boot,17)*32+bytesPerSector-1)/bytesPerSector);
            if(boot[13]==0||total<(uint)reserved+fats*fatSectors+roots)throw new IOException("Invalid FAT data range.");
            if((total-(uint)reserved-(uint)fats*fatSectors-roots)/boot[13]<4085)
                return new(stream,"FAT12",[],512,bytesPerSector,false);
        }
        long[] fatOffsets=Enumerable.Range(0,fats).Select(i=>((long)reserved+i*fatSectors)*bytesPerSector).ToArray();bool unclean=false;
        foreach(long at in fatOffsets){byte[] b=Read(stream,at,512);unclean|=fat32?(U32(b,4)&0x08000000)==0:(U16(b,2)&0x8000)==0;}
        return new(stream,fat32?"FAT32":"FAT16",fatOffsets,512,bytesPerSector,unclean);
    }
    // Restores the USA trailers before examining attributes. The returned
    // buffer is deprotected until Protect is called before the sector write.
    static int NtfsFlags(byte[] raw,int sector)
    {
        if(!raw.AsSpan(0,4).SequenceEqual("FILE"u8))throw new IOException("Invalid NTFS volume record.");
        int usa=U16(raw,4),count=U16(raw,6);
        if(count!=raw.Length/sector+1||usa<8||usa+count*2>raw.Length)throw new IOException("Invalid NTFS record fixup range.");
        ushort sequence=U16(raw,usa);
        for(int i=1;i<count;i++){int end=i*sector-2;if(U16(raw,end)!=sequence)throw new IOException("Torn NTFS volume record.");W16(raw,end,U16(raw,usa+i*2));}
        int used=checked((int)U32(raw,24)),at=U16(raw,20);
        if(used<at||used>raw.Length)throw new IOException("Invalid NTFS record length.");
        for(int budget=0;budget<128&&at+8<=used;budget++)
        {
            uint type=U32(raw,at);if(type==0xffffffff)break;int length=checked((int)U32(raw,at+4));
            if(length<24||length>used-at)throw new IOException("Invalid NTFS attribute bounds.");
            if(type==0x70)
            {
                int offset=U16(raw,at+20);uint size=U32(raw,at+16);
                if(raw[at+8]!=0||raw[at+9]!=0||size<12||offset<24||offset+size>length)throw new IOException("Invalid NTFS volume-information attribute.");
                return at+offset+10;
            }
            at+=length;
        }
        throw new IOException("The NTFS volume-information attribute is missing.");
    }
    static void Protect(byte[] raw,int sector)
    {
        int usa=U16(raw,4),count=U16(raw,6);ushort sequence=unchecked((ushort)(U16(raw,usa)+1));if(sequence==0)sequence=1;W16(raw,usa,sequence);
        for(int i=1;i<count;i++){int end=i*sector-2;W16(raw,usa+i*2,U16(raw,end));W16(raw,end,sequence);}
    }
    public void Set(bool dirty)
    {
        var patches=new List<(long Offset,byte[] Bytes)>();
        foreach(long at in offsets)
        {
            byte[] bytes=Read(stream,at,recordBytes);
            if(kind=="NTFS")
            {int flags=NtfsFlags(bytes,sectorBytes);ushort value=U16(bytes,flags);W16(bytes,flags,dirty?(ushort)(value|1):(ushort)(value&~1));Protect(bytes,sectorBytes);}
            else if(kind=="exFAT")
            {ushort value=U16(bytes,106);W16(bytes,106,dirty?(ushort)(value|2):(ushort)(value&~2));}
            else if(kind=="FAT32")
            {uint value=U32(bytes,4);BinaryPrimitives.WriteUInt32LittleEndian(bytes.AsSpan(4),dirty?value&~0x08000000u:value|0x08000000u);}
            else {ushort value=U16(bytes,2);W16(bytes,2,dirty?(ushort)(value&~0x8000):(ushort)(value|0x8000));}
            patches.Add((at,bytes));
        }
        foreach(var patch in patches){stream.Position=patch.Offset;stream.Write(patch.Bytes);}
        stream.Flush();
    }
}