XFESystem

using System.Numerics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Runtime.Intrinsics;
using XFESystem.Desktop;

namespace XFESystem.Platform;

/// <summary>Present RGB24 CPU pixels to an actual scanout. A failed copy must throw.</summary>
public interface IFramebufferPresenter
{
    int Width { get; }
    int Height { get; }
    void Present(ReadOnlySpan<int> rgbPixels, DesktopBounds nativeDamage);
}

/// <summary>A mapped, little-endian physical framebuffer. Ownership/mapping belongs to the kernel.</summary>
public readonly record struct FramebufferSurface(nuint Address, nuint MappedBytes, int Width, int Height,
    int PitchBytes, int BitsPerPixel, uint RedMask, uint GreenMask, uint BlueMask);

/// <summary>Writes the supplied framebuffer address directly, respecting pitch and RGB masks.</summary>
public sealed unsafe class MemoryFramebufferPresenter : IFramebufferPresenter
{
    private readonly FramebufferSurface surface;
    private readonly int bytesPerPixel;
    private readonly bool packedBgr32, packedRgb32;
    public int Width => surface.Width;
    public int Height => surface.Height;

    public MemoryFramebufferPresenter(FramebufferSurface surface)
    {
        if (surface.Address == 0 || surface.Width <= 0 || surface.Height <= 0 || surface.BitsPerPixel is not (16 or 24 or 32))
            throw new ArgumentException("A mapped 16, 24 or 32 bit framebuffer is required.", nameof(surface));
        bytesPerPixel = surface.BitsPerPixel / 8;
        if (surface.PitchBytes < (long)surface.Width * bytesPerPixel)
            throw new ArgumentException("Framebuffer pitch is smaller than a visible row.", nameof(surface));
        ulong required = checked((ulong)(surface.Height - 1) * (uint)surface.PitchBytes + (ulong)surface.Width * (uint)bytesPerPixel);
        if (surface.PitchBytes < (long)surface.Width * bytesPerPixel || required > surface.MappedBytes ||
            required > nuint.MaxValue - surface.Address || !ValidMask(surface.RedMask, surface.BitsPerPixel) ||
            !ValidMask(surface.GreenMask, surface.BitsPerPixel) || !ValidMask(surface.BlueMask, surface.BitsPerPixel) ||
            (surface.RedMask & surface.GreenMask) != 0 || (surface.RedMask & surface.BlueMask) != 0 || (surface.GreenMask & surface.BlueMask) != 0)
            throw new ArgumentException("Framebuffer geometry, mapping extent or RGB masks are invalid.", nameof(surface));
        this.surface = surface;
        packedBgr32 = surface.BitsPerPixel == 32 && surface.RedMask == 0x00FF0000 && surface.GreenMask == 0x0000FF00 && surface.BlueMask == 0x000000FF;
        packedRgb32 = surface.BitsPerPixel == 32 && surface.RedMask == 0x000000FF && surface.GreenMask == 0x0000FF00 && surface.BlueMask == 0x00FF0000;
    }

    public void Present(ReadOnlySpan<int> rgbPixels, DesktopBounds nativeDamage)
    {
        if ((long)Width * Height > rgbPixels.Length) throw new ArgumentException("The CPU pixel buffer is incomplete.", nameof(rgbPixels));
        var damage = PixelGeometry.Intersect(nativeDamage, new(0, 0, Width, Height));
        for (int y = damage.Y; y < damage.Y + damage.Height; y++)
        {
            byte* output = (byte*)surface.Address + (nuint)y * (nuint)surface.PitchBytes + (nuint)damage.X * (uint)bytesPerPixel;
            if (packedBgr32 || packedRgb32)
            {
                PresentPacked32(rgbPixels.Slice(y * Width + damage.X, damage.Width), output, packedRgb32);
                continue;
            }
            for (int x = damage.X; x < damage.X + damage.Width; x++)
            {
                uint rgb = unchecked((uint)rgbPixels[y * Width + x]);
                uint packed = Channel((byte)(rgb >> 16), surface.RedMask) | Channel((byte)(rgb >> 8), surface.GreenMask) | Channel((byte)rgb, surface.BlueMask);
                for (int b = 0; b < bytesPerPixel; b++) *output++ = (byte)(packed >> (b * 8));
            }
        }
        // Limine may map the framebuffer write-combining. Complete the CPU
        // stores before publishing this copy as complete; this is not vblank.
        if (damage.Width > 0 && damage.Height > 0) System.Threading.Thread.MemoryBarrier();
    }

    private static void PresentPacked32(ReadOnlySpan<int> pixels, byte* output, bool swapRedBlue)
    {
        int x = 0;
        // Use 128-bit vectors only: the independent kernel saves the complete
        // SSE/FPU state, but does not yet switch AVX upper-register state.
        // All reads/stores stay inside the clipped row, including unaligned pitch.
        if (Vector128.IsHardwareAccelerated)
        {
            ref byte input = ref Unsafe.As<int, byte>(ref MemoryMarshal.GetReference(pixels));
            var rgb = Vector128.Create(0x00FFFFFFu);
            var green = Vector128.Create(0x0000FF00u);
            var lowByte = Vector128.Create(0x000000FFu);
            for (; x <= pixels.Length - Vector128<uint>.Count; x += Vector128<uint>.Count)
            {
                var value = Unsafe.ReadUnaligned<Vector128<uint>>(ref Unsafe.Add(ref input, x * 4));
                value = swapRedBlue
                    ? ((value & lowByte) << 16) | (value & green) | ((value >> 16) & lowByte)
                    : value & rgb;
                Unsafe.WriteUnaligned(output + x * 4, value);
            }
        }
        for (; x < pixels.Length; x++)
        {
            uint value = unchecked((uint)pixels[x]);
            uint packed = swapRedBlue ? ((value & 255) << 16) | (value & 0xFF00) | ((value >> 16) & 255) : value & 0xFFFFFF;
            Unsafe.WriteUnaligned(output + x * 4, packed);
        }
    }

    private static bool ValidMask(uint mask, int bits)
    {
        if (mask == 0 || (bits < 32 && (mask >> bits) != 0)) return false;
        uint aligned = mask >> BitOperations.TrailingZeroCount(mask);
        return (aligned & (aligned + 1)) == 0;
    }
    private static uint Channel(byte value, uint mask)
    {
        int shift = BitOperations.TrailingZeroCount(mask);
        uint range = mask >> shift;
        return (uint)(((ulong)value * range + 127) / 255) << shift;
    }
}

/// <summary>Real host-memory scanout for rendering tests and offscreen consumers.</summary>
public sealed class ArrayFramebufferPresenter : IFramebufferPresenter
{
    private readonly int[] pixels;
    public int Width { get; }
    public int Height { get; }
    public ReadOnlyMemory<int> Pixels => pixels;
    public long Presentations { get; private set; }
    public ArrayFramebufferPresenter(int width, int height)
    {
        PixelGeometry.ValidateSize(width, height);
        Width = width; Height = height; pixels = new int[checked(width * height)];
    }
    public void Present(ReadOnlySpan<int> rgbPixels, DesktopBounds nativeDamage)
    {
        if (rgbPixels.Length < pixels.Length) throw new ArgumentException("The CPU pixel buffer is incomplete.", nameof(rgbPixels));
        var damage = PixelGeometry.Intersect(nativeDamage, new(0, 0, Width, Height));
        for (int y = damage.Y; y < damage.Y + damage.Height; y++)
            rgbPixels.Slice(y * Width + damage.X, damage.Width).CopyTo(pixels.AsSpan(y * Width + damage.X, damage.Width));
        Presentations++;
    }
}

internal static class PixelGeometry
{
    // 8K fits; reject overflowing/accidental enormous managed allocations.
    internal const int MaximumPixels = 33_554_432;
    internal static void ValidateSize(int width, int height)
    {
        if (width <= 0 || height <= 0 || (long)width * height > MaximumPixels) throw new ArgumentOutOfRangeException(nameof(width));
    }
    internal static DesktopBounds Intersect(DesktopBounds a, DesktopBounds b)
    {
        long x = Math.Max(a.X, b.X), y = Math.Max(a.Y, b.Y);
        long right = Math.Min((long)a.X + Math.Max(0, a.Width), (long)b.X + Math.Max(0, b.Width));
        long bottom = Math.Min((long)a.Y + Math.Max(0, a.Height), (long)b.Y + Math.Max(0, b.Height));
        if (right <= x || bottom <= y) return default;
        return new((int)x, (int)y, (int)(right - x), (int)(bottom - y));
    }
}