poe2-bot/src/Poe2Trade.Navigation/DesktopDuplication.cs

using System.Runtime.InteropServices;
using OpenCvSharp;
using Serilog;
using SharpGen.Runtime;
using Vortice.Direct3D;
using Vortice.Direct3D11;
using Vortice.DXGI;
using Region = Poe2Trade.Core.Region;

namespace Poe2Trade.Navigation;

public sealed class DesktopDuplication : IScreenCapture
{
    private readonly ID3D11Device _device;
    private readonly ID3D11DeviceContext _context;
    private IDXGIOutputDuplication? _duplication;
    private ID3D11Texture2D? _staging;
    private int _stagingW, _stagingH;
    private bool _needsRecreate;

    public DesktopDuplication()
    {
        D3D11.D3D11CreateDevice(
            null,
            DriverType.Hardware,
            DeviceCreationFlags.BgraSupport,
            [FeatureLevel.Level_11_0],
            out _device!,
            out _context!).CheckError();

        CreateDuplication();
    }

    private void CreateDuplication()
    {
        using var dxgiDevice = _device.QueryInterface<IDXGIDevice>();
        using var adapter = dxgiDevice.GetAdapter();
        adapter.EnumOutputs(0, out var output);
        using var _ = output;
        using var output1 = output.QueryInterface<IDXGIOutput1>();
        _duplication = output1.DuplicateOutput(_device);
        _needsRecreate = false;
        Log.Debug("DXGI Desktop Duplication created");
    }

    public unsafe Mat? CaptureRegion(Region region)
    {
        if (_duplication == null) return null;

        if (_needsRecreate)
        {
            try
            {
                _duplication.Dispose();
                CreateDuplication();
            }
            catch (Exception ex)
            {
                Log.Warning(ex, "Failed to recreate DXGI duplication");
                return null;
            }
        }

        IDXGIResource? resource = null;
        try
        {
            _duplication.AcquireNextFrame(100, out _, out resource);
        }
        catch (SharpGenException ex)
        {
            if (ex.ResultCode == Vortice.DXGI.ResultCode.AccessLost)
                _needsRecreate = true;
            // WaitTimeout is normal when screen hasn't changed
            return null;
        }

        try
        {
            using var srcTexture = resource!.QueryInterface<ID3D11Texture2D>();
            EnsureStaging(region.Width, region.Height);

            // Copy only the region we need from the desktop texture
            _context.CopySubresourceRegion(
                _staging!, 0, 0, 0, 0,
                srcTexture, 0,
                new Vortice.Mathematics.Box(region.X, region.Y, 0,
                    region.X + region.Width, region.Y + region.Height, 1));

            var mapped = _context.Map(_staging!, 0, MapMode.Read);
            try
            {
                var mat = new Mat(region.Height, region.Width, MatType.CV_8UC4);
                var rowBytes = region.Width * 4;

                for (var row = 0; row < region.Height; row++)
                {
                    Buffer.MemoryCopy(
                        (void*)(mapped.DataPointer + row * mapped.RowPitch),
                        (void*)mat.Ptr(row),
                        rowBytes, rowBytes);
                }

                // BGRA → BGR
                var bgr = new Mat();
                Cv2.CvtColor(mat, bgr, ColorConversionCodes.BGRA2BGR);
                mat.Dispose();
                return bgr;
            }
            finally
            {
                _context.Unmap(_staging!, 0);
            }
        }
        finally
        {
            resource?.Dispose();
            _duplication!.ReleaseFrame();
        }
    }

    private void EnsureStaging(int w, int h)
    {
        if (_staging != null && _stagingW == w && _stagingH == h) return;
        _staging?.Dispose();

        _staging = _device.CreateTexture2D(new Texture2DDescription
        {
            Width = (uint)w,
            Height = (uint)h,
            MipLevels = 1,
            ArraySize = 1,
            Format = Format.B8G8R8A8_UNorm,
            SampleDescription = new SampleDescription(1, 0),
            Usage = ResourceUsage.Staging,
            CPUAccessFlags = CpuAccessFlags.Read,
        });
        _stagingW = w;
        _stagingH = h;
    }

    public void Dispose()
    {
        _staging?.Dispose();
        _duplication?.Dispose();
        _context?.Dispose();
        _device?.Dispose();
    }
}