XFESystem

using System.Buffers.Binary;
using System.Text;
using XFESystem.Core;

namespace XFESystem.Storage;

public interface IPhysicalDisk
{
    string Name { get; }
    ulong SectorCount { get; }
    int LogicalSectorBytes => 512;
    bool CanWrite { get; }
    void Read(ulong lba, Span<byte> bytes);
    void Write(ulong lba, ReadOnlySpan<byte> bytes);
    void Flush();
    void AuthorizePartition(ulong start, ulong count);
}
public readonly record struct DiskPartition(int Number, ulong Start, ulong Count, bool DataPartition, string Name);

public static class DiskPartitions
{
    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 ulong Sectors(IPhysicalDisk disk)=>disk.SectorCount/(ulong)(disk.LogicalSectorBytes/512);
    static DiskPartition[] Normalize(IEnumerable<DiskPartition> partitions,IPhysicalDisk disk)
    {ulong unit=(ulong)(disk.LogicalSectorBytes/512);return partitions.Select(p=>p with{Start=checked(p.Start*unit),Count=checked(p.Count*unit)}).ToArray();}
    static byte[] Read(IPhysicalDisk disk,ulong lba,int count=1)
    {
        if(lba>=Sectors(disk)||(ulong)count>Sectors(disk)-lba)throw new IOException("Partition metadata leaves its disk.");
        byte[] bytes=new byte[checked(count*disk.LogicalSectorBytes)];disk.Read(lba*(ulong)(disk.LogicalSectorBytes/512),bytes);return bytes;
    }
    public static DiskPartition[] Read(IPhysicalDisk disk)
    {
        if(disk.LogicalSectorBytes is not(512 or 4096)||disk.SectorCount%(ulong)(disk.LogicalSectorBytes/512)!=0)throw new IOException("Unsupported disk sector geometry.");
        byte[] mbr=Read(disk,0);if(mbr[510]!=85||mbr[511]!=170)return [];
        var result=new List<DiskPartition>();
        if(Enumerable.Range(0,4).Any(i=>mbr[446+i*16+4]==0xee))return Normalize(Gpt(disk),disk);
        var containers=new List<(ulong Start,ulong Count)>();
        for(int i=0;i<4;i++)
        {
            int at=446+i*16;byte type=mbr[at+4];ulong start=U32(mbr,at+8),count=U32(mbr,at+12);
            if(type==0||count==0)continue;Bounds(disk,start,count);
            if(type is 5 or 15 or 133){containers.Add((start,count));continue;}
            result.Add(new(i+1,start,count,type is 7 or 6 or 11 or 12 or 14,""));
        }
        Validate(result.Concat(containers.Select(c=>new DiskPartition(0,c.Start,c.Count,false,""))).ToList());
        int number=5;var ebrSectors=new List<ulong>();
        foreach(var container in containers)
        {
            ulong current=container.Start;var seen=new HashSet<ulong>();
            for(int limit=0;limit<128;limit++)
            {
                if(!seen.Add(current)||current<container.Start||current-container.Start>=container.Count)throw new IOException("Cyclic or invalid extended partition chain.");
                ebrSectors.Add(current);
                byte[] ebr=Read(disk,current);if(ebr[510]!=85||ebr[511]!=170)throw new IOException("Invalid extended partition signature.");
                byte type=ebr[450];ulong relative=U32(ebr,454),count=U32(ebr,458);
                if(type!=0&&count!=0)
                {
                    if(relative==0||relative>ulong.MaxValue-current)throw new IOException("Invalid logical partition offset.");
                    ulong start=current+relative;Bounds(disk,start,count);
                    if(start<container.Start||start-container.Start>=container.Count||count>container.Count-(start-container.Start))throw new IOException("Logical partition leaves its container.");
                    result.Add(new(number++,start,count,type is 7 or 6 or 11 or 12 or 14,""));
                }
                byte linkType=ebr[466];ulong link=U32(ebr,470);
                if(linkType==0&&link==0)break;
                if(linkType is not(5 or 15 or 133)||link==0||link>=container.Count||limit==127)throw new IOException("Invalid extended partition link.");
                current=container.Start+link;
            }
        }
        Validate(result);
        if(ebrSectors.Any(lba=>result.Any(p=>lba>=p.Start&&lba-p.Start<p.Count)))throw new IOException("An extended partition record overlaps file data.");
        return Normalize(result,disk);
    }
    static DiskPartition[] Gpt(IPhysicalDisk disk)
    {
        ulong sectors=Sectors(disk);byte[] header=Read(disk,1);Header(header,1,sectors-1,sectors);
        uint count=U32(header,80),size=U32(header,84);
        if(count==0||count>128||size<128||size>1024||(size&7)!=0)throw new IOException("Unsupported GPT entry geometry.");
        int bytes=checked((int)(count*size)),tableSectors=(bytes+disk.LogicalSectorBytes-1)/disk.LogicalSectorBytes;
        byte[] table=Read(disk,U64(header,72),tableSectors);
        if(UsbWorkspaceFormat.Crc(table.AsSpan(0,bytes))!=U32(header,88))throw new IOException("GPT partition CRC mismatch.");
        byte[] backup=Read(disk,sectors-1);Header(backup,sectors-1,1,sectors);
        if(!header.AsSpan(56,16).SequenceEqual(backup.AsSpan(56,16))||U32(backup,80)!=count||U32(backup,84)!=size||U32(backup,88)!=U32(header,88)||
           !Read(disk,U64(backup,72),tableSectors).AsSpan(0,bytes).SequenceEqual(table.AsSpan(0,bytes)))throw new IOException("GPT primary and backup disagree.");
        ulong first=U64(header,40),last=U64(header,48);
        ulong primaryTable=U64(header,72),backupTable=U64(backup,72);
        if(first<2+(ulong)tableSectors||last<first||last>=sectors-1||first!=U64(backup,40)||last!=U64(backup,48)||
           primaryTable<2||primaryTable>=first||(ulong)tableSectors>first-primaryTable||
           backupTable<=last||backupTable>=sectors-1||(ulong)tableSectors>sectors-1-backupTable)throw new IOException("Invalid GPT usable range or table placement.");
        var result=new List<DiskPartition>();var ids=new HashSet<Guid>();
        for(int i=0;i<count;i++)
        {
            int at=checked(i*(int)size);var type=new Guid(table.AsSpan(at,16));if(type==Guid.Empty)continue;
            var id=new Guid(table.AsSpan(at+16,16));ulong start=U64(table,at+32),end=U64(table,at+40);
            if(id==Guid.Empty||!ids.Add(id)||start<first||end<start||end>last)throw new IOException("Invalid GPT partition identity or range.");
            bool data=type==new Guid("ebd0a0a2-b9e5-4433-87c0-68b6b72699c7")&&(U64(table,at+48)&(1UL<<60))==0;
            result.Add(new(i+1,start,end-start+1,data,Encoding.Unicode.GetString(table,at+56,72).TrimEnd('\0')));
        }
        Validate(result);return result.ToArray();
    }
    static void Header(byte[] header,ulong current,ulong other,ulong sectors)
    {
        uint size=U32(header,12);if(!header.AsSpan(0,8).SequenceEqual("EFI PART"u8)||U32(header,8)!=65536||size<92||size>header.Length||U32(header,20)!=0||U64(header,24)!=current||U64(header,32)!=other||other>=sectors)throw new IOException("Invalid GPT header.");
        uint crc=U32(header,16);byte[] copy=header.AsSpan(0,(int)size).ToArray();copy.AsSpan(16,4).Clear();
        if(UsbWorkspaceFormat.Crc(copy)!=crc)throw new IOException("GPT header CRC mismatch.");
    }
    static void Bounds(IPhysicalDisk disk,ulong start,ulong count)
    {if(start==0||count==0||start>=Sectors(disk)||count>Sectors(disk)-start)throw new IOException("Partition leaves its disk.");}
    static void Validate(List<DiskPartition> parts)
    {
        var sorted=parts.OrderBy(p=>p.Start).ToArray();
        for(int i=1;i<sorted.Length;i++)if(sorted[i].Start-sorted[i-1].Start<sorted[i-1].Count)throw new IOException("Overlapping partitions were refused.");
    }
}