using System.Buffers.Binary;
using System.Text;

namespace MatrixImageEditor;

public sealed class Fat16Disk : IDisposable
{
    private readonly FileStream stream;
    private readonly long fatOffset;
    private readonly long rootOffset;
    private readonly int rootEntries;
    private readonly int fatBytes;
    private readonly long dataOffset;
    private readonly int clusterBytes;

    public Fat16Disk(string path, bool writable)
    {
        stream = new FileStream(path, writable ? FileMode.Open : FileMode.Open, writable ? FileAccess.ReadWrite : FileAccess.Read, FileShare.Read, 1024 * 1024, FileOptions.RandomAccess);
        var mbr = ReadAt(0, 512);
        if (mbr[510] != 0x55 || mbr[511] != 0xAA) throw new InvalidDataException("文件没有有效的 MBR 分区表");
        var partitionLba = BinaryPrimitives.ReadUInt32LittleEndian(mbr.AsSpan(454, 4));
        var partitionSectors = BinaryPrimitives.ReadUInt32LittleEndian(mbr.AsSpan(458, 4));
        var partitionType = mbr[450];
        // WalnutPi labels this BPB-compatible FAT16 volume as 0x0C in the MBR.
        if (partitionLba == 0 || partitionSectors == 0 || partitionType is not (0x04 or 0x06 or 0x0B or 0x0C or 0x0E)) throw new InvalidDataException("镜像第一分区不是受支持的 FAT 启动分区");
        var partitionOffset = partitionLba * 512L;
        if (partitionOffset + partitionSectors * 512L > stream.Length) throw new InvalidDataException("镜像分区表超出文件边界");
        var bpb = ReadAt(partitionOffset, 512);
        var bytesPerSector = BinaryPrimitives.ReadUInt16LittleEndian(bpb.AsSpan(11, 2));
        var sectorsPerCluster = bpb[13];
        var reserved = BinaryPrimitives.ReadUInt16LittleEndian(bpb.AsSpan(14, 2));
        var fatCount = bpb[16];
        rootEntries = BinaryPrimitives.ReadUInt16LittleEndian(bpb.AsSpan(17, 2));
        var fatSectors = BinaryPrimitives.ReadUInt16LittleEndian(bpb.AsSpan(22, 2));
        if (bytesPerSector != 512 || sectorsPerCluster == 0 || fatSectors == 0 || rootEntries == 0) throw new InvalidDataException("镜像第一分区不是受支持的 FAT16");
        clusterBytes = bytesPerSector * sectorsPerCluster;
        fatBytes = fatSectors * bytesPerSector;
        fatOffset = partitionOffset + reserved * bytesPerSector;
        rootOffset = fatOffset + fatCount * fatBytes;
        var rootSectors = (rootEntries * 32L + bytesPerSector - 1) / bytesPerSector;
        dataOffset = rootOffset + rootSectors * bytesPerSector;
    }

    public byte[] ReadFile(string name)
    {
        var entry = Find(name);
        var fat = ReadAt(fatOffset, fatBytes);
        var result = new byte[entry.Size];
        var written = 0;
        foreach (var cluster in Chain(entry.FirstCluster, fat))
        {
            var count = Math.Min(clusterBytes, result.Length - written);
            ReadAt(ClusterOffset(cluster), result.AsSpan(written, count));
            written += count;
            if (written == result.Length) break;
        }
        if (written != result.Length) throw new InvalidDataException("镜像配置文件被截断");
        return result;
    }

    public void WriteFileRange(string name, int fileOffset, byte[] data)
    {
        var entry = Find(name);
        if (fileOffset < 0 || fileOffset + data.Length > entry.Size) throw new ArgumentOutOfRangeException(nameof(fileOffset));
        var fat = ReadAt(fatOffset, fatBytes);
        var chain = Chain(entry.FirstCluster, fat).ToArray();
        var remaining = data.Length;
        var sourceOffset = 0;
        var position = fileOffset;
        while (remaining > 0)
        {
            var chainIndex = position / clusterBytes;
            var within = position % clusterBytes;
            if (chainIndex >= chain.Length) throw new InvalidDataException("镜像配置文件簇链不足");
            var count = Math.Min(remaining, clusterBytes - within);
            stream.Position = ClusterOffset(chain[chainIndex]) + within;
            stream.Write(data, sourceOffset, count);
            position += count; sourceOffset += count; remaining -= count;
        }
        stream.Flush(flushToDisk: true);
    }

    private (ushort FirstCluster, int Size) Find(string name)
    {
        var expected = Name83(name);
        var root = ReadAt(rootOffset, rootEntries * 32);
        for (var index = 0; index < rootEntries; index++)
        {
            var entry = root.AsSpan(index * 32, 32);
            if (entry[0] == 0x00) break;
            if (entry[0] == 0xE5 || entry[11] == 0x0F) continue;
            if (entry[..11].SequenceEqual(expected))
            {
                return (BinaryPrimitives.ReadUInt16LittleEndian(entry.Slice(26, 2)), BinaryPrimitives.ReadInt32LittleEndian(entry.Slice(28, 4)));
            }
        }
        throw new InvalidDataException($"镜像缺少 {name}");
    }

    private IEnumerable<ushort> Chain(ushort first, byte[] fat)
    {
        var seen = new HashSet<ushort>();
        var cluster = first;
        while (cluster is >= 2 and < 0xFFF8)
        {
            if (!seen.Add(cluster)) throw new InvalidDataException("FAT16 簇链循环");
            yield return cluster;
            cluster = BinaryPrimitives.ReadUInt16LittleEndian(fat.AsSpan(cluster * 2, 2));
        }
    }

    private long ClusterOffset(ushort cluster) => dataOffset + (cluster - 2L) * clusterBytes;
    private byte[] ReadAt(long offset, int count) { var value = new byte[count]; ReadAt(offset, value); return value; }
    private void ReadAt(long offset, Span<byte> destination) { stream.Position = offset; stream.ReadExactly(destination); }

    private static byte[] Name83(string name)
    {
        var parts = name.ToUpperInvariant().Split('.', 2);
        if (parts[0].Length is < 1 or > 8 || (parts.Length == 2 && parts[1].Length > 3)) throw new ArgumentException("文件名不是 8.3 格式");
        return Encoding.ASCII.GetBytes(parts[0].PadRight(8) + (parts.Length == 2 ? parts[1] : "").PadRight(3));
    }

    public void Dispose() => stream.Dispose();
}