LlamaCasty/Services/ScreenCaptureFrameSource.cs

using System;
using System.Runtime.InteropServices;
using System.Runtime.InteropServices.WindowsRuntime;
using System.Threading.Tasks;
using Windows.Graphics;
using Windows.Graphics.Capture;
using Windows.Graphics.DirectX;
using Windows.Graphics.DirectX.Direct3D11;
using Windows.Graphics.Imaging;
using ytLive.Helpers;

namespace ytLive.Services;

/// <summary>
/// A live screen capture for one target key ("monitor:&lt;n&gt;" or
/// "window:&lt;hwnd&gt;"). Owns the GraphicsCaptureItem, a free-threaded
/// Direct3D11CaptureFramePool and the capture session; frames are converted on
/// the capture worker thread from the GPU surface to a CPU BGRA VideoFrame.
/// Known OS limits: DRM content captures as black frames; capture pauses while
/// the app is minimized (the frame pool simply stops delivering frames).
/// </summary>
public sealed class ScreenCaptureFrameSource : IScreenCaptureSource
{
    private readonly object _gate = new();
    private GraphicsCaptureItem _item;
    private Direct3D11CaptureFramePool? _framePool;
    private GraphicsCaptureSession? _session;
    private SizeInt32 _poolSize;
    private bool _started;
    private bool _framePending;
    private DateTime _lastErrorLog = DateTime.MinValue;

    // The composition master frame (see ai.md "Resolution tiers"): the background
    // is an input layer, so we never hold a CPU frame bigger than the master.
    private const int MaxBackgroundWidth = 1920;
    private const int MaxBackgroundHeight = 1080;

    // A failing conversion must not re-flood the log at frame rate.
    private static readonly TimeSpan ErrorLogThrottle = TimeSpan.FromSeconds(5);

    public string Key { get; }
    public event Action<VideoFrame>? FrameAvailable;

    internal ScreenCaptureFrameSource(string key, GraphicsCaptureItem item)
    {
        Key = key;
        _item = item;
    }

    /// <summary>Wraps an item chosen through the OS GraphicsCapturePicker.</summary>
    public static ScreenCaptureFrameSource CreateForPicker(string key, GraphicsCaptureItem item)
        => new(key, item);

    public static ScreenCaptureFrameSource CreateForMonitor(int monitorIndex)
    {
        var hmonitor = Win32FullScreenDetector.GetMonitorHandle(monitorIndex);
        if (hmonitor == IntPtr.Zero)
            throw new InvalidOperationException($"Monitor {monitorIndex} is not connected");
        var item = CaptureInterop.CreateForMonitor(hmonitor);
        return new ScreenCaptureFrameSource($"monitor:{monitorIndex}", item);
    }

    public static ScreenCaptureFrameSource CreateForWindow(IntPtr hwnd)
    {
        var item = CaptureInterop.CreateForWindow(hwnd);
        return new ScreenCaptureFrameSource($"window:0x{hwnd.ToInt64():X}", item);
    }

    public Task StartAsync()
    {
        lock (_gate)
        {
            if (_started) return Task.CompletedTask;
            var item = _item;
            var device = Direct3D11Helper.CreateDevice();
            var framePool = Direct3D11CaptureFramePool.CreateFreeThreaded(
                device, DirectXPixelFormat.B8G8R8A8UIntNormalized, 2, item.Size);
            _poolSize = item.Size;
            var session = framePool.CreateCaptureSession(item);
            framePool.FrameArrived += OnFrameArrived;
            session.StartCapture();
            _framePool = framePool;
            _session = session;
            _started = true;
        }
        return Task.CompletedTask;
    }

    public Task StopAsync()
    {
        lock (_gate)
        {
            _started = false;
            _framePending = false;
            if (_framePool != null)
                _framePool.FrameArrived -= OnFrameArrived;
            _session?.Dispose();
            _session = null;
            _framePool?.Dispose();
            _framePool = null;
        }
        return Task.CompletedTask;
    }

    private void OnFrameArrived(Direct3D11CaptureFramePool sender, object args)
    {
        var frame = sender.TryGetNextFrame();
        if (frame == null) return;
        lock (_gate)
        {
            if (!_started)
            {
                frame.Dispose();
                return;
            }
        }

        if (frame.ContentSize.Width != _poolSize.Width || frame.ContentSize.Height != _poolSize.Height)
        {
            sender.Recreate(Direct3D11Helper.CreateDevice(),
                DirectXPixelFormat.B8G8R8A8UIntNormalized, 2, frame.ContentSize);
            _poolSize = frame.ContentSize;
        }

        if (_framePending)
        {
            frame.Dispose();
            return;
        }
        _framePending = true;
        _ = ProcessFrameAsync(frame);
    }

    private async Task ProcessFrameAsync(Direct3D11CaptureFrame frame)
    {
        try
        {
            using (frame)
            using (var softwareBitmap = await SoftwareBitmap.CreateCopyFromSurfaceAsync(
                frame.Surface, BitmapAlphaMode.Ignore))
            {
                FrameAvailable?.Invoke(CopyToVideoFrame(softwareBitmap));
            }
        }
        catch (Exception ex)
        {
            var now = DateTime.UtcNow;
            if (now - _lastErrorLog >= ErrorLogThrottle)
            {
                _lastErrorLog = now;
                AppLog.Write($"ScreenCaptureFrameSource: frame conversion failed: {ex.Message}");
            }
        }
        finally
        {
            _framePending = false;
        }
    }

    private static VideoFrame CopyToVideoFrame(SoftwareBitmap bitmap)
    {
        var sw = bitmap.PixelWidth;
        var sh = bitmap.PixelHeight;
        using var buffer = bitmap.LockBuffer(BitmapBufferAccessMode.Read);
        using var reference = buffer.CreateReference();
        if (!WindowsRuntimeMarshal.TryGetDataUnsafe(reference, out var data, out var capacity))
            throw new InvalidOperationException("Could not access the frame buffer.");

        var srcStride = sw * 4;
        var count = (int)Math.Min(capacity, (uint)(sh * srcStride));

        // The master frame is 1920×1080; a larger monitor is scaled down here so
        // the CPU never holds a frame above the master (the rendered layer fills
        // the frame with UniformToFill regardless).
        if (sw > MaxBackgroundWidth || sh > MaxBackgroundHeight)
        {
            var scale = Math.Min(MaxBackgroundWidth / (double)sw, MaxBackgroundHeight / (double)sh);
            var dw = Math.Max(1, (int)(sw * scale));
            var dh = Math.Max(1, (int)(sh * scale));
            // DWM delivers an opaque surface (alpha 255); bilinear keeps it 255.
            return new VideoFrame(dw, dh, DownscaleBgra(data, sw, sh, srcStride, dw, dh)) { IsOpaque = true };
        }

        var pixels = new byte[count];
        Marshal.Copy(data, pixels, 0, pixels.Length);
        return new VideoFrame(sw, sh, pixels) { IsOpaque = true };
    }

    // Bilinear downscale to the master frame. Reads each source row pair through
    // Marshal.Copy (no unsafe), writing tightly packed BGRA output.
    private static byte[] DownscaleBgra(IntPtr src, int sw, int sh, int srcStride, int dw, int dh)
    {
        var row0 = new byte[srcStride];
        var row1 = new byte[srcStride];
        var dst = new byte[dw * dh * 4];
        var xs = sw / (double)dw;
        var ys = sh / (double)dh;

        for (var y = 0; y < dh; y++)
        {
            var sy = Math.Min(sh - 1, (int)(y * ys));
            var sy1 = Math.Min(sh - 1, sy + 1);
            var fy = (y * ys) - sy;
            Marshal.Copy(IntPtr.Add(src, sy * srcStride), row0, 0, srcStride);
            Marshal.Copy(IntPtr.Add(src, sy1 * srcStride), row1, 0, srcStride);

            var dRow = y * dw * 4;
            for (var x = 0; x < dw; x++)
            {
                var sx = Math.Min(sw - 1, (int)(x * xs));
                var sx1 = Math.Min(sw - 1, sx + 1);
                var fx = (x * xs) - sx;
                for (var c = 0; c < 4; c++)
                {
                    var i0 = sx * 4 + c;
                    var i1 = sx1 * 4 + c;
                    var top = row0[i0] + (row0[i1] - row0[i0]) * fx;
                    var bottom = row1[i0] + (row1[i1] - row1[i0]) * fx;
                    dst[dRow + x * 4 + c] = (byte)(top + (bottom - top) * fy);
                }
            }
        }
        return dst;
    }
}