LightlessClient/LightlessSync/Services/TextureCompression/TexFileHelper.cs

using System.Runtime.InteropServices;
using Lumina.Data.Files;
using OtterTex;

namespace LightlessSync.Services.TextureCompression;

// base taken from penumbra mostly
internal static class TexFileHelper
{
    private const int HeaderSize = 80;
    private const int MaxMipLevels = 13;

    public static ScratchImage Load(string path)
    {
        using var stream = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.ReadWrite | FileShare.Delete);
        return Load(stream);
    }

    public static ScratchImage Load(Stream stream)
    {
        using var reader = new BinaryReader(stream, System.Text.Encoding.UTF8, leaveOpen: true);
        var header = ReadHeader(reader);
        var meta = CreateMeta(header);
        meta.MipLevels = ComputeMipCount(stream.Length, header, meta);
        if (meta.MipLevels == 0)
        {
            throw new InvalidOperationException("TEX file does not contain a valid mip chain.");
        }

        var scratch = ScratchImage.Initialize(meta);
        ReadPixelData(reader, scratch);
        return scratch;
    }

    public static void Save(string path, ScratchImage image)
    {
        var header = BuildHeader(image);
        if (header.Format == TexFile.TextureFormat.Unknown)
        {
            throw new InvalidOperationException($"Unable to export TEX file with unsupported format {image.Meta.Format}.");
        }

        var mode = File.Exists(path) ? FileMode.Truncate : FileMode.CreateNew;
        using var stream = new FileStream(path, mode, FileAccess.Write, FileShare.Read);
        using var writer = new BinaryWriter(stream);
        WriteHeader(writer, header);
        writer.Write(image.Pixels);
        GC.KeepAlive(image);
    }

    private static TexFile.TexHeader ReadHeader(BinaryReader reader)
    {
        Span<byte> buffer = stackalloc byte[HeaderSize];
        var read = reader.Read(buffer);
        if (read != HeaderSize)
        {
            throw new EndOfStreamException($"Incomplete TEX header: expected {HeaderSize} bytes, read {read} bytes.");
        }

        return MemoryMarshal.Read<TexFile.TexHeader>(buffer);
    }

    private static TexMeta CreateMeta(in TexFile.TexHeader header)
    {
        var meta = new TexMeta
        {
            Width = header.Width,
            Height = header.Height,
            Depth = Math.Max(header.Depth, (ushort)1),
            ArraySize = 1,
            MipLevels = header.MipCount,
            Format = header.Format.ToDxgi(),
            Dimension = header.Type.ToDimension(),
            MiscFlags = header.Type.HasFlag(TexFile.Attribute.TextureTypeCube) ? D3DResourceMiscFlags.TextureCube : 0,
            MiscFlags2 = 0,
        };

        if (meta.Format == DXGIFormat.Unknown)
        {
            throw new InvalidOperationException($"TEX format {header.Format} cannot be mapped to DXGI.");
        }

        if (meta.Dimension == TexDimension.Unknown)
        {
            throw new InvalidOperationException($"Unrecognised TEX dimension attribute {header.Type}.");
        }

        return meta;
    }

    private static unsafe int ComputeMipCount(long totalLength, in TexFile.TexHeader header, in TexMeta meta)
    {
        var width = Math.Max(meta.Width, 1);
        var height = Math.Max(meta.Height, 1);
        var minSide = meta.Format.IsCompressed() ? 4 : 1;
        var bitsPerPixel = meta.Format.BitsPerPixel();

        var expectedOffset = HeaderSize;
        var remaining = totalLength - HeaderSize;

        for (var level = 0; level < MaxMipLevels; level++)
        {
            var declaredOffset = header.OffsetToSurface[level];
            if (declaredOffset == 0)
            {
                return level;
            }

            if (declaredOffset != expectedOffset || remaining <= 0)
            {
                return level;
            }

            var mipSize = (int)((long)width * height * bitsPerPixel / 8);
            if (mipSize > remaining)
            {
                return level;
            }

            expectedOffset += mipSize;
            remaining -= mipSize;

            if (width <= minSide && height <= minSide)
            {
                return level + 1;
            }

            width = Math.Max(width / 2, minSide);
            height = Math.Max(height / 2, minSide);
        }

        return MaxMipLevels;
    }

    private static unsafe void ReadPixelData(BinaryReader reader, ScratchImage image)
    {
        fixed (byte* destination = image.Pixels)
        {
            var span = new Span<byte>(destination, image.Pixels.Length);
            var read = reader.Read(span);
            if (read < span.Length)
            {
                throw new InvalidDataException($"TEX pixel buffer is truncated (read {read} of {span.Length} bytes).");
            }
        }
    }

    private static TexFile.TexHeader BuildHeader(ScratchImage image)
    {
        var meta = image.Meta;
        var header = new TexFile.TexHeader
        {
            Width = (ushort)meta.Width,
            Height = (ushort)meta.Height,
            Depth = (ushort)Math.Max(meta.Depth, 1),
            MipCount = (byte)Math.Min(meta.MipLevels, MaxMipLevels),
            Format = meta.Format.ToTex(),
            Type = meta.Dimension switch
            {
                _ when meta.IsCubeMap => TexFile.Attribute.TextureTypeCube,
                TexDimension.Tex1D => TexFile.Attribute.TextureType1D,
                TexDimension.Tex2D => TexFile.Attribute.TextureType2D,
                TexDimension.Tex3D => TexFile.Attribute.TextureType3D,
                _ => 0,
            },
        };

        PopulateOffsets(ref header, image);
        return header;
    }

    private static unsafe void PopulateOffsets(ref TexFile.TexHeader header, ScratchImage image)
    {
        var index = 0;
        fixed (byte* basePtr = image.Pixels)
        {
            foreach (var mip in image.Images)
            {
                if (index >= MaxMipLevels)
                {
                    break;
                }

                var byteOffset = (byte*)mip.Pixels - basePtr;
                header.OffsetToSurface[index++] = HeaderSize + (uint)byteOffset;
            }
        }

        while (index < MaxMipLevels)
        {
            header.OffsetToSurface[index++] = 0;
        }

        header.LodOffset[0] = 0;
        header.LodOffset[1] = (byte)Math.Min(header.MipCount - 1, 1);
        header.LodOffset[2] = (byte)Math.Min(header.MipCount - 1, 2);
    }

    private static unsafe void WriteHeader(BinaryWriter writer, in TexFile.TexHeader header)
    {
        writer.Write((uint)header.Type);
        writer.Write((uint)header.Format);
        writer.Write(header.Width);
        writer.Write(header.Height);
        writer.Write(header.Depth);
        writer.Write((byte)(header.MipCount | (header.MipUnknownFlag ? 0x80 : 0)));
        writer.Write(header.ArraySize);
        writer.Write(header.LodOffset[0]);
        writer.Write(header.LodOffset[1]);
        writer.Write(header.LodOffset[2]);
        for (var i = 0; i < MaxMipLevels; i++)
        {
            writer.Write(header.OffsetToSurface[i]);
        }
    }

    private static TexDimension ToDimension(this TexFile.Attribute attribute)
        => (attribute & TexFile.Attribute.TextureTypeMask) switch
        {
            TexFile.Attribute.TextureType1D => TexDimension.Tex1D,
            TexFile.Attribute.TextureType2D => TexDimension.Tex2D,
            TexFile.Attribute.TextureType3D => TexDimension.Tex3D,
            _ => TexDimension.Unknown,
        };

    private static DXGIFormat ToDxgi(this TexFile.TextureFormat format)
        => format switch
        {
            TexFile.TextureFormat.L8 => DXGIFormat.R8UNorm,
            TexFile.TextureFormat.A8 => DXGIFormat.A8UNorm,
            TexFile.TextureFormat.B4G4R4A4 => DXGIFormat.B4G4R4A4UNorm,
            TexFile.TextureFormat.B5G5R5A1 => DXGIFormat.B5G5R5A1UNorm,
            TexFile.TextureFormat.B8G8R8A8 => DXGIFormat.B8G8R8A8UNorm,
            TexFile.TextureFormat.B8G8R8X8 => DXGIFormat.B8G8R8X8UNorm,
            TexFile.TextureFormat.R32F => DXGIFormat.R32Float,
            TexFile.TextureFormat.R16G16F => DXGIFormat.R16G16Float,
            TexFile.TextureFormat.R32G32F => DXGIFormat.R32G32Float,
            TexFile.TextureFormat.R16G16B16A16F => DXGIFormat.R16G16B16A16Float,
            TexFile.TextureFormat.R32G32B32A32F => DXGIFormat.R32G32B32A32Float,
            TexFile.TextureFormat.BC1 => DXGIFormat.BC1UNorm,
            TexFile.TextureFormat.BC2 => DXGIFormat.BC2UNorm,
            TexFile.TextureFormat.BC3 => DXGIFormat.BC3UNorm,
            (TexFile.TextureFormat)0x6120 => DXGIFormat.BC4UNorm,
            TexFile.TextureFormat.BC5 => DXGIFormat.BC5UNorm,
            (TexFile.TextureFormat)0x6330 => DXGIFormat.BC6HSF16,
            TexFile.TextureFormat.BC7 => DXGIFormat.BC7UNorm,
            TexFile.TextureFormat.D16 => DXGIFormat.R16G16B16A16Typeless,
            TexFile.TextureFormat.D24S8 => DXGIFormat.R24G8Typeless,
            TexFile.TextureFormat.Shadow16 => DXGIFormat.R16Typeless,
            TexFile.TextureFormat.Shadow24 => DXGIFormat.R24G8Typeless,
            _ => DXGIFormat.Unknown,
        };

    private static TexFile.TextureFormat ToTex(this DXGIFormat format)
        => format switch
        {
            DXGIFormat.R8UNorm => TexFile.TextureFormat.L8,
            DXGIFormat.A8UNorm => TexFile.TextureFormat.A8,
            DXGIFormat.B4G4R4A4UNorm => TexFile.TextureFormat.B4G4R4A4,
            DXGIFormat.B5G5R5A1UNorm => TexFile.TextureFormat.B5G5R5A1,
            DXGIFormat.B8G8R8A8UNorm => TexFile.TextureFormat.B8G8R8A8,
            DXGIFormat.B8G8R8X8UNorm => TexFile.TextureFormat.B8G8R8X8,
            DXGIFormat.R32Float => TexFile.TextureFormat.R32F,
            DXGIFormat.R16G16Float => TexFile.TextureFormat.R16G16F,
            DXGIFormat.R32G32Float => TexFile.TextureFormat.R32G32F,
            DXGIFormat.R16G16B16A16Float => TexFile.TextureFormat.R16G16B16A16F,
            DXGIFormat.R32G32B32A32Float => TexFile.TextureFormat.R32G32B32A32F,
            DXGIFormat.BC1UNorm => TexFile.TextureFormat.BC1,
            DXGIFormat.BC2UNorm => TexFile.TextureFormat.BC2,
            DXGIFormat.BC3UNorm => TexFile.TextureFormat.BC3,
            DXGIFormat.BC4UNorm => (TexFile.TextureFormat)0x6120,
            DXGIFormat.BC5UNorm => TexFile.TextureFormat.BC5,
            DXGIFormat.BC6HSF16 => (TexFile.TextureFormat)0x6330,
            DXGIFormat.BC7UNorm => TexFile.TextureFormat.BC7,
            DXGIFormat.R16G16B16A16Typeless => TexFile.TextureFormat.D16,
            DXGIFormat.R24G8Typeless => TexFile.TextureFormat.D24S8,
            DXGIFormat.R16Typeless => TexFile.TextureFormat.Shadow16,
            _ => TexFile.TextureFormat.Unknown,
        };
}