serverRoom/device-agent/Services/H264ScreenCaptureService.cs

using System.Runtime.InteropServices;
using DeviceAgent.Models;
using SharpDX;
using SharpDX.Direct3D11;
using SharpDX.DXGI;
using SharpDX.MediaFoundation;
using Device = SharpDX.Direct3D11.Device;
using Resource = SharpDX.DXGI.Resource;

namespace DeviceAgent.Services;

/// <summary>
/// H.264 屏幕捕获服务 - 使用 DXGI Desktop Duplication + Media Foundation H.264 编码
/// </summary>
public class H264ScreenCaptureService : IDisposable
{
    private readonly ILogger<H264ScreenCaptureService> _logger;
    private Device? _device;
    private OutputDuplication? _duplicatedOutput;
    private Texture2D? _stagingTexture;
    private SinkWriter? _sinkWriter;
    private int _videoStreamIndex;
    private int _frameWidth;
    private int _frameHeight;
    private int _fps;
    private int _bitrate;
    private long _frameIndex;
    private bool _isInitialized;
    private readonly object _lock = new();
    private MemoryStream? _outputStream;
    private byte[]? _lastEncodedFrame;
    private StreamQualityProfile _currentProfile = StreamQualityProfile.Low;
    private int _targetWidth;
    private int _targetHeight;

    public H264ScreenCaptureService(ILogger<H264ScreenCaptureService> logger)
    {
        _logger = logger;
    }

    public bool Initialize(int targetWidth = 1280, int targetHeight = 720, int fps = 15, int bitrate = 2000000)
    {
        lock (_lock)
        {
            try
            {
                if (_isInitialized) return true;

                _targetWidth = targetWidth;
                _targetHeight = targetHeight;
                _fps = fps;
                _bitrate = bitrate;

                // 初始化 Media Foundation
                MediaManager.Startup();

                // 创建 D3D11 设备
                _device = new Device(SharpDX.Direct3D.DriverType.Hardware,
                    DeviceCreationFlags.BgraSupport | DeviceCreationFlags.VideoSupport);

                // 获取 DXGI 输出
                using var dxgiDevice = _device.QueryInterface<SharpDX.DXGI.Device>();
                using var adapter = dxgiDevice.Adapter;
                using var output = adapter.GetOutput(0);
                using var output1 = output.QueryInterface<Output1>();

                // 获取屏幕尺寸
                var outputDesc = output.Description;
                _frameWidth = Math.Min(targetWidth, outputDesc.DesktopBounds.Right - outputDesc.DesktopBounds.Left);
                _frameHeight = Math.Min(targetHeight, outputDesc.DesktopBounds.Bottom - outputDesc.DesktopBounds.Top);

                // 创建桌面复制
                _duplicatedOutput = output1.DuplicateOutput(_device);

                // 创建暂存纹理用于 CPU 读取
                var textureDesc = new Texture2DDescription
                {
                    Width = _frameWidth,
                    Height = _frameHeight,
                    MipLevels = 1,
                    ArraySize = 1,
                    Format = Format.B8G8R8A8_UNorm,
                    SampleDescription = new SampleDescription(1, 0),
                    Usage = ResourceUsage.Staging,
                    CpuAccessFlags = CpuAccessFlags.Read,
                    BindFlags = BindFlags.None
                };
                _stagingTexture = new Texture2D(_device, textureDesc);

                // 初始化 H.264 编码器
                InitializeEncoder(fps, bitrate);

                _isInitialized = true;
                _logger.LogInformation("H.264 屏幕捕获服务初始化成功: {Width}x{Height}, {Fps}fps, {Bitrate}bps",
                    _frameWidth, _frameHeight, fps, bitrate);
                return true;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "初始化 H.264 屏幕捕获服务失败");
                Cleanup();
                return false;
            }
        }
    }

    /// <summary>
    /// 设置质量档位（动态切换）
    /// </summary>
    public bool SetQuality(StreamQualityProfile profile)
    {
        lock (_lock)
        {
            try
            {
                _logger.LogInformation("切换质量档位: {Profile}", profile);
                _currentProfile = profile;
                
                // 清理现有资源
                Cleanup();
                
                // 使用新参数重新初始化
                return Initialize(profile.Width, profile.Height, profile.Fps, profile.Bitrate);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "切换质量档位失败");
                return false;
            }
        }
    }

    private void InitializeEncoder(int fps, int bitrate)
    {
        _outputStream = new MemoryStream();

        // 创建字节流
        var byteStream = new ByteStream(_outputStream);

        // 创建 Sink Writer 属性
        using var attributes = new MediaAttributes();
        attributes.Set(SinkWriterAttributeKeys.ReadwriteEnableHardwareTransforms, 1);

        // 创建 Sink Writer
        _sinkWriter = MediaFactory.CreateSinkWriterFromURL(null, byteStream, attributes);

        // 设置输出媒体类型 (H.264)
        using var outputType = new MediaType();
        outputType.Set(MediaTypeAttributeKeys.MajorType, MediaTypeGuids.Video);
        outputType.Set(MediaTypeAttributeKeys.Subtype, VideoFormatGuids.H264);
        outputType.Set(MediaTypeAttributeKeys.AvgBitrate, bitrate);
        outputType.Set(MediaTypeAttributeKeys.InterlaceMode, (int)VideoInterlaceMode.Progressive);
        outputType.Set(MediaTypeAttributeKeys.FrameSize, PackSize(_frameWidth, _frameHeight));
        outputType.Set(MediaTypeAttributeKeys.FrameRate, PackSize(fps, 1));
        outputType.Set(MediaTypeAttributeKeys.PixelAspectRatio, PackSize(1, 1));

        _sinkWriter.AddStream(outputType, out _videoStreamIndex);

        // 设置输入媒体类型 (BGRA)
        using var inputType = new MediaType();
        inputType.Set(MediaTypeAttributeKeys.MajorType, MediaTypeGuids.Video);
        inputType.Set(MediaTypeAttributeKeys.Subtype, VideoFormatGuids.Argb32);
        inputType.Set(MediaTypeAttributeKeys.InterlaceMode, (int)VideoInterlaceMode.Progressive);
        inputType.Set(MediaTypeAttributeKeys.FrameSize, PackSize(_frameWidth, _frameHeight));
        inputType.Set(MediaTypeAttributeKeys.FrameRate, PackSize(fps, 1));
        inputType.Set(MediaTypeAttributeKeys.PixelAspectRatio, PackSize(1, 1));

        _sinkWriter.SetInputMediaType(_videoStreamIndex, inputType, null);
        _sinkWriter.BeginWriting();
    }

    private static long PackSize(int width, int height)
    {
        return ((long)width << 32) | (uint)height;
    }

    /// <summary>
    /// 捕获并编码一帧
    /// </summary>
    public byte[]? CaptureFrame()
    {
        lock (_lock)
        {
            if (!_isInitialized || _duplicatedOutput == null || _device == null)
                return null;

            try
            {
                // 尝试获取下一帧
                var result = _duplicatedOutput.TryAcquireNextFrame(100,
                    out var frameInfo, out var desktopResource);

                if (result.Failure)
                {
                    return _lastEncodedFrame; // 返回上一帧
                }

                try
                {
                    using var desktopTexture = desktopResource.QueryInterface<Texture2D>();

                    // 复制到暂存纹理
                    _device.ImmediateContext.CopyResource(desktopTexture, _stagingTexture);

                    // 读取像素数据
                    var dataBox = _device.ImmediateContext.MapSubresource(
                        _stagingTexture, 0, MapMode.Read, SharpDX.Direct3D11.MapFlags.None);

                    try
                    {
                        // 编码帧
                        var encodedFrame = EncodeFrame(dataBox.DataPointer, dataBox.RowPitch);
                        if (encodedFrame != null && encodedFrame.Length > 0)
                        {
                            _lastEncodedFrame = encodedFrame;
                        }
                    }
                    finally
                    {
                        _device.ImmediateContext.UnmapSubresource(_stagingTexture, 0);
                    }
                }
                finally
                {
                    desktopResource?.Dispose();
                    _duplicatedOutput.ReleaseFrame();
                }

                return _lastEncodedFrame;
            }
            catch (SharpDXException ex) when (ex.ResultCode == SharpDX.DXGI.ResultCode.AccessLost)
            {
                _logger.LogWarning("桌面访问丢失，需要重新初始化");
                _isInitialized = false;
                return null;
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "捕获帧失败");
                return _lastEncodedFrame;
            }
        }
    }

    private unsafe byte[]? EncodeFrame(IntPtr dataPointer, int rowPitch)
    {
        if (_sinkWriter == null || _outputStream == null)
            return null;

        try
        {
            var frameSize = _frameWidth * _frameHeight * 4;

            // 创建媒体缓冲区
            var buffer = MediaFactory.CreateMemoryBuffer(frameSize);

            try
            {
                // 锁定缓冲区并复制数据
                var bufferPtr = buffer.Lock(out var maxLength, out var currentLength);
                try
                {
                    // 复制像素数据
                    for (int y = 0; y < _frameHeight; y++)
                    {
                        var srcRow = IntPtr.Add(dataPointer, y * rowPitch);
                        var dstRow = IntPtr.Add(bufferPtr, y * _frameWidth * 4);
                        System.Buffer.MemoryCopy(srcRow.ToPointer(), dstRow.ToPointer(),
                            _frameWidth * 4, _frameWidth * 4);
                    }
                }
                finally
                {
                    buffer.Unlock();
                }

                buffer.CurrentLength = frameSize;

                // 创建样本
                using var sample = MediaFactory.CreateSample();
                sample.AddBuffer(buffer);

                // 设置时间戳
                var duration = 10_000_000L / 15; // 假设 15fps
                sample.SampleTime = _frameIndex * duration;
                sample.SampleDuration = duration;

                // 重置输出流
                _outputStream.SetLength(0);
                _outputStream.Position = 0;

                // 写入样本
                _sinkWriter.WriteSample(_videoStreamIndex, sample);

                _frameIndex++;

                // 返回编码后的数据
                if (_outputStream.Length > 0)
                {
                    return _outputStream.ToArray();
                }
            }
            finally
            {
                buffer?.Dispose();
            }
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "编码帧失败");
        }

        return null;
    }

    private void Cleanup()
    {
        _isInitialized = false;

        try { _sinkWriter?.Dispose(); } catch { }
        try { _stagingTexture?.Dispose(); } catch { }
        try { _duplicatedOutput?.Dispose(); } catch { }
        try { _device?.Dispose(); } catch { }
        try { _outputStream?.Dispose(); } catch { }

        _sinkWriter = null;
        _stagingTexture = null;
        _duplicatedOutput = null;
        _device = null;
        _outputStream = null;
    }

    public void Dispose()
    {
        lock (_lock)
        {
            Cleanup();
            MediaManager.Shutdown();
        }
    }
}