using System.Collections.Concurrent;
using System.Net.WebSockets;

namespace AmtScanner.Api.Services;

/// <summary>
/// 屏幕流代理服务 - 管理到 Agent 的 WebSocket 连接
/// </summary>
public class ScreenStreamProxyService
{
    private readonly ILogger<ScreenStreamProxyService> _logger;
    private readonly ConcurrentDictionary<string, AgentConnection> _agentConnections = new();

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

    /// <summary>
    /// 代理屏幕流 - 从 Agent 获取屏幕流并转发给客户端
    /// </summary>
    public async Task ProxyScreenStreamAsync(
        WebSocket clientWebSocket,
        string agentIp,
        int agentPort,
        CancellationToken cancellationToken)
    {
        ClientWebSocket? agentWebSocket = null;
        
        try
        {
            // 连接到 Agent 的 WebSocket 服务
            agentWebSocket = new ClientWebSocket();
            var agentUri = new Uri($"ws://{agentIp}:{agentPort}/");
            
            _logger.LogInformation("正在连接到 Agent: {Uri}", agentUri);
            
            using var connectCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
            using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, connectCts.Token);
            
            await agentWebSocket.ConnectAsync(agentUri, linkedCts.Token);
            
            _logger.LogInformation("已连接到 Agent: {Ip}:{Port}", agentIp, agentPort);

            // 转发数据
            var buffer = new byte[1024 * 1024]; // 1MB buffer for images
            
            while (agentWebSocket.State == WebSocketState.Open && 
                   clientWebSocket.State == WebSocketState.Open &&
                   !cancellationToken.IsCancellationRequested)
            {
                try
                {
                    var result = await agentWebSocket.ReceiveAsync(
                        new ArraySegment<byte>(buffer), 
                        cancellationToken);

                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        _logger.LogInformation("Agent 关闭了连接");
                        break;
                    }

                    if (result.MessageType == WebSocketMessageType.Binary && result.Count > 0)
                    {
                        // 转发给客户端
                        await clientWebSocket.SendAsync(
                            new ArraySegment<byte>(buffer, 0, result.Count),
                            WebSocketMessageType.Binary,
                            result.EndOfMessage,
                            cancellationToken);
                    }
                }
                catch (WebSocketException ex)
                {
                    _logger.LogWarning("WebSocket 错误: {Message}", ex.Message);
                    break;
                }
            }
        }
        catch (WebSocketException ex)
        {
            _logger.LogError("连接 Agent 失败: {Message}", ex.Message);
            
            // 发送错误消息给客户端
            if (clientWebSocket.State == WebSocketState.Open)
            {
                var errorMsg = System.Text.Encoding.UTF8.GetBytes($"ERROR:无法连接到设备 {agentIp}:{agentPort}");
                await clientWebSocket.SendAsync(
                    new ArraySegment<byte>(errorMsg),
                    WebSocketMessageType.Text,
                    true,
                    cancellationToken);
            }
        }
        catch (OperationCanceledException)
        {
            _logger.LogInformation("代理连接已取消");
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "代理屏幕流时发生错误");
        }
        finally
        {
            if (agentWebSocket != null)
            {
                try
                {
                    if (agentWebSocket.State == WebSocketState.Open)
                    {
                        await agentWebSocket.CloseAsync(
                            WebSocketCloseStatus.NormalClosure, 
                            "Proxy closing", 
                            CancellationToken.None);
                    }
                    agentWebSocket.Dispose();
                }
                catch { }
            }
            
            _logger.LogInformation("代理连接已关闭: {Ip}:{Port}", agentIp, agentPort);
        }
    }
}

/// <summary>
/// Agent 连接信息
/// </summary>
public class AgentConnection
{
    public string Uuid { get; set; } = "";
    public string IpAddress { get; set; } = "";
    public int Port { get; set; }
    public ClientWebSocket? WebSocket { get; set; }
    public DateTime ConnectedAt { get; set; }
}