serverRoom/device-agent/Services/ScreenStreamService.cs

using System.Net.WebSockets;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Options;
using DeviceAgent.Models;

namespace DeviceAgent.Services;

/// <summary>
/// 质量切换请求
/// </summary>
internal class QualityChangeRequest
{
    public string? Quality { get; set; }
}

/// <summary>
/// 屏幕流服务 - 通过 WebSocket 实时推送 H.264 编码的屏幕画面
/// </summary>
public class ScreenStreamService : IDisposable
{
    private readonly ILogger<ScreenStreamService> _logger;
    private readonly ScreenCaptureService _screenCaptureService;
    private readonly H264ScreenCaptureService _h264CaptureService;
    private readonly AgentConfig _config;
    private WebApplication? _app;
    private readonly List<WebSocket> _clients = new();
    private readonly object _clientsLock = new();
    private CancellationTokenSource? _cts;
    private Task? _streamTask;
    private bool _isRunning;
    private bool _useH264;
    private StreamQualityProfile _currentQuality = StreamQualityProfile.Low;

    public ScreenStreamService(
        ILogger<ScreenStreamService> logger,
        ScreenCaptureService screenCaptureService,
        H264ScreenCaptureService h264CaptureService,
        IOptions<AgentConfig> config)
    {
        _logger = logger;
        _screenCaptureService = screenCaptureService;
        _h264CaptureService = h264CaptureService;
        _config = config.Value;
    }

    public async Task StartAsync(CancellationToken cancellationToken)
    {
        if (!_config.ScreenStreamEnabled)
        {
            _logger.LogInformation("屏幕流服务已禁用");
            return;
        }

        try
        {
            _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);

            // 使用低质量档位初始化（默认监控墙模式）
            _currentQuality = StreamQualityProfile.Low;
            
            // 尝试初始化 H.264 编码
            if (_config.UseH264Encoding)
            {
                _useH264 = _h264CaptureService.Initialize(
                    _currentQuality.Width,
                    _currentQuality.Height,
                    _currentQuality.Fps,
                    _currentQuality.Bitrate);

                if (_useH264)
                {
                    _logger.LogInformation("使用 H.264 编码模式，初始质量: {Quality}", _currentQuality);
                }
                else
                {
                    _logger.LogWarning("H.264 初始化失败，回退到 JPEG 模式");
                }
            }

            var builder = WebApplication.CreateSlimBuilder();
            builder.WebHost.ConfigureKestrel(options =>
            {
                options.ListenAnyIP(_config.ScreenStreamPort);
            });
            builder.Logging.ClearProviders();

            _app = builder.Build();
            _app.UseWebSockets();

            _app.Map("/", async context =>
            {
                if (context.WebSockets.IsWebSocketRequest)
                {
                    var webSocket = await context.WebSockets.AcceptWebSocketAsync();
                    await HandleWebSocketAsync(webSocket, _cts.Token);
                }
                else
                {
                    context.Response.StatusCode = 200;
                    var mode = _useH264 ? "H.264" : "JPEG";
                    await context.Response.WriteAsync($"Screen Stream ({mode}) - Clients: {_clients.Count}");
                }
            });

            // 提供流信息端点
            _app.Map("/info", async context =>
            {
                context.Response.ContentType = "application/json";
                await context.Response.WriteAsJsonAsync(new
                {
                    mode = _useH264 ? "h264" : "jpeg",
                    width = _currentQuality.Width,
                    height = _currentQuality.Height,
                    fps = _currentQuality.Fps,
                    bitrate = _currentQuality.Bitrate,
                    quality = _currentQuality.Level.ToString(),
                    clients = _clients.Count
                });
            });
            
            // 质量控制端点
            _app.Map("/quality", async context =>
            {
                if (context.Request.Method == "POST")
                {
                    try
                    {
                        var body = await context.Request.ReadFromJsonAsync<QualityChangeRequest>();
                        if (body != null)
                        {
                            var newQuality = body.Quality?.ToLower() == "high" 
                                ? StreamQualityProfile.High 
                                : StreamQualityProfile.Low;
                            
                            if (SetQuality(newQuality))
                            {
                                await context.Response.WriteAsJsonAsync(new { success = true, quality = newQuality.Level.ToString() });
                            }
                            else
                            {
                                context.Response.StatusCode = 500;
                                await context.Response.WriteAsJsonAsync(new { success = false, error = "切换质量失败" });
                            }
                        }
                    }
                    catch (Exception ex)
                    {
                        _logger.LogError(ex, "处理质量切换请求失败");
                        context.Response.StatusCode = 500;
                        await context.Response.WriteAsJsonAsync(new { success = false, error = ex.Message });
                    }
                }
                else
                {
                    await context.Response.WriteAsJsonAsync(new { quality = _currentQuality.Level.ToString() });
                }
            });

            _isRunning = true;
            _logger.LogInformation("屏幕流服务已启动，端口: {Port}, 模式: {Mode}",
                _config.ScreenStreamPort, _useH264 ? "H.264" : "JPEG");

            _streamTask = StreamScreenAsync(_cts.Token);
            await _app.RunAsync(_cts.Token);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "启动屏幕流服务失败");
        }
    }

    private async Task HandleWebSocketAsync(WebSocket webSocket, CancellationToken ct)
    {
        try
        {
            lock (_clientsLock) { _clients.Add(webSocket); }
            _logger.LogInformation("客户端连接，当前: {Count}, 模式: {Mode}",
                _clients.Count, _useH264 ? "H.264" : "JPEG");

            // 发送初始化消息告知客户端编码模式
            var initMsg = System.Text.Encoding.UTF8.GetBytes(
                System.Text.Json.JsonSerializer.Serialize(new
                {
                    type = "init",
                    mode = _useH264 ? "h264" : "jpeg",
                    width = _currentQuality.Width,
                    height = _currentQuality.Height,
                    fps = _currentQuality.Fps,
                    quality = _currentQuality.Level.ToString()
                }));
            await webSocket.SendAsync(new ArraySegment<byte>(initMsg),
                WebSocketMessageType.Text, true, ct);

            var buffer = new byte[1024];
            while (webSocket.State == WebSocketState.Open && !ct.IsCancellationRequested)
            {
                try
                {
                    var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), ct);
                    if (result.MessageType == WebSocketMessageType.Close) break;
                }
                catch { break; }
            }
        }
        finally
        {
            lock (_clientsLock) { _clients.Remove(webSocket); }
            _logger.LogInformation("客户端断开，当前: {Count}", _clients.Count);
            try { webSocket.Dispose(); } catch { }
        }
    }

    private async Task StreamScreenAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested && _isRunning)
        {
            try
            {
                List<WebSocket> clients;
                lock (_clientsLock) { clients = _clients.ToList(); }

                // 按需推流：只在有客户端连接时才采集编码
                if (clients.Count > 0)
                {
                    byte[]? frameData;

                    if (_useH264)
                    {
                        // 使用 H.264 编码
                        frameData = _h264CaptureService.CaptureFrame();
                    }
                    else
                    {
                        // 回退到 JPEG
                        frameData = _screenCaptureService.CaptureScreen(
                            _config.ScreenStreamQuality, _currentQuality.Width);
                    }

                    if (frameData != null && frameData.Length > 0)
                    {
                        var tasks = clients
                            .Where(ws => ws.State == WebSocketState.Open)
                            .Select(ws => SendFrameAsync(ws, frameData, ct));
                        await Task.WhenAll(tasks);
                    }
                }
                
                // 根据当前质量档位动态调整帧间隔
                var interval = TimeSpan.FromMilliseconds(1000.0 / _currentQuality.Fps);
                await Task.Delay(interval, ct);
            }
            catch (OperationCanceledException) { break; }
            catch (Exception ex)
            {
                _logger.LogError(ex, "推送屏幕失败");
                await Task.Delay(1000, ct);
            }
        }
    }
    
    /// <summary>
    /// 设置流质量（公开方法，供 SignalingClientService 调用）
    /// </summary>
    public bool SetQuality(StreamQualityProfile profile)
    {
        try
        {
            if (_currentQuality.Level == profile.Level)
            {
                return true; // 已是目标质量，无需切换
            }

            _logger.LogInformation("切换流质量: {OldQuality} -> {NewQuality}", _currentQuality, profile);
            _currentQuality = profile;
            
            if (_useH264)
            {
                return _h264CaptureService.SetQuality(profile);
            }
            
            return true;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "切换流质量失败");
            return false;
        }
    }

    private async Task SendFrameAsync(WebSocket ws, byte[] frame, CancellationToken ct)
    {
        try
        {
            await ws.SendAsync(new ArraySegment<byte>(frame), WebSocketMessageType.Binary, true, ct);
        }
        catch { }
    }

    public async Task StopAsync()
    {
        _isRunning = false;
        _cts?.Cancel();

        List<WebSocket> clients;
        lock (_clientsLock) { clients = _clients.ToList(); _clients.Clear(); }

        foreach (var ws in clients)
        {
            try { ws.Dispose(); } catch { }
        }

        if (_app != null)
        {
            await _app.StopAsync();
            await _app.DisposeAsync();
        }

        _logger.LogInformation("屏幕流服务已停止");
    }

    public void Dispose()
    {
        _cts?.Dispose();
        _h264CaptureService?.Dispose();
    }
}