XFESystem

using System.Buffers.Binary;
using DiscUtils;
using DiscUtils.ExFat;
using DiscUtils.Fat;
using DiscUtils.Ntfs;
using XFESystem.Core;
using XFESystem.Storage;

namespace XFESystem.StorageFixtures;

// Host-only fixture code. FileStream is never used by the guest storage path.
public sealed class FixtureDisk : IPhysicalDisk, IDisposable
{
    readonly FileStream stream;
    ulong allowedStart,allowedCount;
    public string Name {get;}
    public bool CanWrite {get;set;}=true;
    public ulong SectorCount=>(ulong)stream.Length/512;
    public int LogicalSectorBytes {get;}
    public ulong? FailWriteAt;
    public int? FailAfterWrites;
    public int Writes,Flushes;
    public List<(ulong Lba,int Bytes)> Written=[];
    public FixtureDisk(string path,string name,bool create=false,long length=128L*1024*1024,int logicalSectorBytes=512)
    {
        string full=Path.GetFullPath(path);
        if(!full.EndsWith(".img",StringComparison.OrdinalIgnoreCase)||full.StartsWith("\\\\")||full.Contains("PhysicalDrive",StringComparison.OrdinalIgnoreCase))throw new ArgumentException("A regular .img fixture file is required.");
        if(logicalSectorBytes is not(512 or 4096))throw new ArgumentException("Fixture sector size.");
        LogicalSectorBytes=logicalSectorBytes;Name=name;stream=new(full,create?FileMode.CreateNew:FileMode.Open,FileAccess.ReadWrite,FileShare.Read);if(create)stream.SetLength(length);
    }
    public void Read(ulong lba,Span<byte> bytes){Bounds(lba,bytes.Length);stream.Position=checked((long)lba*512);stream.ReadExactly(bytes);}
    void Bounds(ulong lba,int bytes){if(bytes%512!=0||lba>=SectorCount||(ulong)bytes/512>SectorCount-lba)throw new IOException("Fixture disk bounds.");}
    public void Write(ulong lba,ReadOnlySpan<byte> bytes)
    {
        Bounds(lba,bytes.Length);
        if(!CanWrite||lba<allowedStart||lba-allowedStart>=allowedCount||(ulong)bytes.Length/512>allowedCount-(lba-allowedStart)||FailWriteAt==lba||FailAfterWrites is{} threshold&&Writes>=threshold)throw new IOException("Fixture write rejected.");
        stream.Position=(long)lba*512;stream.Write(bytes);Written.Add((lba,bytes.Length));Writes++;
    }
    public void Raw(ulong lba,ReadOnlySpan<byte> bytes){Bounds(lba,bytes.Length);stream.Position=(long)lba*512;stream.Write(bytes);}
    public void AuthorizePartition(ulong start,ulong count){if(start>=SectorCount||count==0||count>SectorCount-start)throw new IOException();allowedStart=start;allowedCount=count;}
    public void Flush(){stream.Flush(true);Flushes++;}
    public void Dispose()=>stream.Dispose();
    static void U16(byte[] b,int at,ushort n)=>BinaryPrimitives.WriteUInt16LittleEndian(b.AsSpan(at),n);
    static void U32(byte[] b,int at,uint n)=>BinaryPrimitives.WriteUInt32LittleEndian(b.AsSpan(at),n);
    static void U64(byte[] b,int at,ulong n)=>BinaryPrimitives.WriteUInt64LittleEndian(b.AsSpan(at),n);
    public DiskPartition Gpt()
    {
        ulong unit=(ulong)(LogicalSectorBytes/512),sectors=SectorCount/unit,tableSectors=16384UL/(ulong)LogicalSectorBytes;
        ulong first=2+tableSectors,last=sectors-tableSectors-2,start=2048/unit;
        byte[] mbr=new byte[512];mbr[450]=0xee;U32(mbr,454,1);U32(mbr,458,(uint)(sectors-1));mbr[510]=85;mbr[511]=170;Raw(0,mbr);
        byte[] entries=new byte[16384];new Guid("ebd0a0a2-b9e5-4433-87c0-68b6b72699c7").TryWriteBytes(entries);Guid.NewGuid().TryWriteBytes(entries.AsSpan(16));
        U64(entries,32,start);U64(entries,40,last);Raw(2*unit,entries);Raw((sectors-1-tableSectors)*unit,entries);
        byte[] id=Guid.NewGuid().ToByteArray();
        foreach(bool backup in new[]{false,true})
        {
            byte[] header=new byte[512];"EFI PART"u8.CopyTo(header);U32(header,8,65536);U32(header,12,92);
            U64(header,24,backup?sectors-1:1);U64(header,32,backup?1:sectors-1);U64(header,40,first);U64(header,48,last);id.CopyTo(header,56);
            U64(header,72,backup?sectors-1-tableSectors:2);U32(header,80,128);U32(header,84,128);U32(header,88,UsbWorkspaceFormat.Crc(entries));U32(header,16,UsbWorkspaceFormat.Crc(header.AsSpan(0,92)));Raw((backup?sectors-1:1)*unit,header);
        }
        return new(1,start*unit,(last-start+1)*unit,true,"");
    }
    public DiskPartition[] ExtendedMbr()
    {
        byte[] mbr=new byte[512];mbr[450]=15;U32(mbr,454,2048);U32(mbr,458,(uint)SectorCount-2048);mbr[510]=85;mbr[511]=170;Raw(0,mbr);
        byte[] first=new byte[512];first[450]=7;U32(first,454,2048);U32(first,458,245760);first[466]=15;U32(first,470,247808);U32(first,474,(uint)SectorCount-249856);first[510]=85;first[511]=170;Raw(2048,first);
        byte[] second=new byte[512];second[450]=12;U32(second,454,2048);U32(second,458,(uint)SectorCount-251904);second[510]=85;second[511]=170;Raw(249856,second);
        return [new(5,4096,245760,true,""),new(6,251904,SectorCount-251904,true,"")];
    }
    public void Format(DiskPartition partition,string format,byte[] program)
    {
        using var content=new PartitionStream(this,partition,true);long sectors=(long)partition.Count;
        DiscFileSystem fs=format switch {
            "NTFS"=>NtfsFileSystem.Format(content,Name,Geometry.FromCapacity((long)SectorCount*512),(long)partition.Start,sectors),
            "exFAT"=>FormatExFat(content,Name),
            _=>FatFileSystem.FormatPartition(content,Name,Geometry.FromCapacity((long)SectorCount*512),(int)partition.Start,(int)partition.Count,0)};
        using(fs)
        {
            fs.CreateDirectory("Programs");fs.CreateDirectory("目录");
            Put(fs,"Readme.txt",System.Text.Encoding.UTF8.GetBytes(Name+"_MEDIA_READ_OK"));Put(fs,"EditMe.txt","before"u8.ToArray());Put(fs,"DeleteMe.txt","delete proof"u8.ToArray());
            Put(fs,"Programs\\External.dll",program);Put(fs,"目录\\中文.txt",System.Text.Encoding.UTF8.GetBytes("真实磁盘中文文件"));
            byte[] large=new byte[196608];for(int i=0;i<large.Length;i++)large[i]=(byte)(i*37+11);Put(fs,"Pattern.bin",large);
            if(format=="NTFS")for(int i=0;i<180;i++)Put(fs,"DirectoryIndex"+i.ToString("D3")+".txt","index"u8.ToArray());
        }
        Flush();Writes=Flushes=0;Written.Clear();
    }
    static void Put(DiscFileSystem fs,string path,byte[] bytes){using var file=fs.OpenFile(path,FileMode.Create,FileAccess.Write);file.Write(bytes);}
    static DiscFileSystem FormatExFat(Stream content,string name)
    {using(DiscUtils.ExFat.Internal.Filesystem.ExFatPathFilesystem.Format(content,new DiscUtils.ExFat.Internal.ExFatFormatOptions(),name)){}return new ExFatFileSystem(content);}
}