YZWater3/src/YZWater.Core/Services/HubClient.cs

using System.Buffers;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using Serilog;
using YZWater.Core.Models;

namespace YZWater.Core.Services;

/// <summary>
/// 中控客户端 - 连接到服务端，接收数据，发送命令，自动重连
/// </summary>
public class HubClient : IDisposable
{
    private ClientWebSocket? _ws;
    private CancellationTokenSource? _cts;
    private Task? _connectionTask;
    private bool _disposed;
    private volatile bool _isConnected;
    private volatile bool _shouldReconnect = true;

    /// <summary>
    /// 服务端地址
    /// </summary>
    public string ServerUrl { get; set; } = "ws://localhost:8765";

    /// <summary>
    /// 认证令牌
    /// </summary>
    public string Token { get; set; } = string.Empty;

    /// <summary>
    /// 是否忽略证书验证（仅用于自签名证书的测试环境）
    /// </summary>
    public bool IgnoreCertificateErrors { get; set; } = false;

    /// <summary>
    /// 是否已连接
    /// </summary>
    public bool IsConnected => _isConnected;

    /// <summary>
    /// 是否有操作锁
    /// </summary>
    public bool HasLock { get; private set; }

    /// <summary>
    /// 当前操作员
    /// </summary>
    public string? CurrentOperator { get; private set; }

    public event Action<PlcDataModel>? DataReceived;
    public event Action<AlarmRecord>? AlarmReceived;
    public event Action<bool, string?>? LockStatusChanged;
    public event Action<bool>? ConnectionChanged;
    public event Action<int>? ClientCountChanged;
    public event Action<bool, string>? CommandResult;

    /// <summary>
    /// 连接到服务端（带自动重连）
    /// </summary>
    public async Task ConnectAsync()
    {
        if (_isConnected) return;

        _cts = new CancellationTokenSource();
        _shouldReconnect = true;

        await ConnectInternalAsync();

        // 启动接收 + 重连循环
        _connectionTask = Task.Run(() => ConnectionLoopAsync(_cts.Token));
    }

    /// <summary>
    /// 断开连接
    /// </summary>
    public async Task DisconnectAsync()
    {
        _shouldReconnect = false;
        _cts?.Cancel();

        if (_ws?.State == WebSocketState.Open)
        {
            try
            {
                await _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "客户端断开", CancellationToken.None);
            }
            catch { /* 断开清理中，忽略关闭异常 */ }
        }

        _isConnected = false;
        HasLock = false;
        CurrentOperator = null;
        ConnectionChanged?.Invoke(false);
        Log.Information("已断开中控服务端连接");
    }

    /// <summary>
    /// 连接 + 接收 + 自动重连主循环
    /// </summary>
    private async Task ConnectionLoopAsync(CancellationToken ct)
    {
        var reconnectDelay = 1000;
        const int maxDelay = 30000;

        while (!ct.IsCancellationRequested && _shouldReconnect)
        {
            try
            {
                if (!_isConnected)
                {
                    await ConnectInternalAsync();
                    reconnectDelay = 1000; // 连接成功，重置延迟
                }

                // 进入接收循环
                await ReceiveLoopAsync(ct);
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (Exception ex)
            {
                Log.Warning(ex, "连接异常");
            }

            // 连接断开，自动重连
            if (!ct.IsCancellationRequested && _shouldReconnect)
            {
                _isConnected = false;
                HasLock = false;
                ConnectionChanged?.Invoke(false);

                Log.Information("将在 {Delay}ms 后重连...", reconnectDelay);
                try
                {
                    await Task.Delay(reconnectDelay, ct);
                }
                catch (OperationCanceledException)
                {
                    break;
                }

                reconnectDelay = Math.Min(reconnectDelay * 2, maxDelay);
            }
        }
    }

    /// <summary>
    /// 内部连接方法
    /// </summary>
    private async Task ConnectInternalAsync()
    {
        _ws?.Dispose();
        _ws = new ClientWebSocket();

        // 如果启用忽略证书验证（自签名证书场景）
        if (IgnoreCertificateErrors && ServerUrl.StartsWith("wss://"))
        {
            _ws.Options.RemoteCertificateValidationCallback = (sender, certificate, chain, sslPolicyErrors) => true;
            Log.Warning("已忽略 TLS 证书验证（仅用于测试环境）");
        }

        try
        {
            await _ws.ConnectAsync(new Uri(ServerUrl), _cts?.Token ?? CancellationToken.None);

            // 发送认证令牌
            if (!string.IsNullOrEmpty(Token))
            {
                await SendMessageAsync(HubMessage.Auth(Token));

                // 等待认证结果
                var authResult = await ReceiveOneMessageAsync(_cts?.Token ?? CancellationToken.None);
                if (authResult?.Type != "auth_result" || authResult.GetPayload<bool>() != true)
                {
                    Log.Warning("中控服务端认证失败");
                    _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "认证失败", CancellationToken.None).Wait(1000);
                    throw new Exception("认证失败");
                }
                Log.Information("中控服务端认证成功");
            }

            _isConnected = true;
            ConnectionChanged?.Invoke(true);
            Log.Information("已连接到中控服务端: {Url}", ServerUrl);
        }
        catch (Exception ex)
        {
            Log.Warning(ex, "连接中控服务端失败: {Url}", ServerUrl);
            _isConnected = false;
            throw;
        }
    }

    /// <summary>
    /// 请求操作锁
    /// </summary>
    public async Task RequestLockAsync(string operatorName)
    {
        await SendMessageAsync(new HubMessage
        {
            Type = "command",
            Payload = JsonSerializer.SerializeToElement(new CommandPayload
            {
                Action = "request_lock",
                Operator = operatorName
            })
        });
    }

    /// <summary>
    /// 释放操作锁
    /// </summary>
    public async Task ReleaseLockAsync()
    {
        await SendMessageAsync(new HubMessage
        {
            Type = "command",
            Payload = JsonSerializer.SerializeToElement(new CommandPayload
            {
                Action = "release_lock",
                Operator = CurrentOperator ?? ""
            })
        });
    }

    /// <summary>
    /// 发送控制命令
    /// </summary>
    public async Task SendCommandAsync(string action, string? target = null, object? value = null, string? operatorName = null)
    {
        await SendMessageAsync(new HubMessage
        {
            Type = "command",
            Payload = JsonSerializer.SerializeToElement(new CommandPayload
            {
                Action = action,
                Target = target,
                Value = value,
                Operator = operatorName ?? CurrentOperator ?? "unknown"
            })
        });
    }

    /// <summary>
    /// 发送消息
    /// </summary>
    private async Task SendMessageAsync(HubMessage message)
    {
        if (_ws?.State != WebSocketState.Open)
        {
            Log.Warning("未连接到服务端，无法发送消息");
            return;
        }

        var bytes = Encoding.UTF8.GetBytes(message.ToJson());
        await _ws.SendAsync(new ArraySegment<byte>(bytes), WebSocketMessageType.Text, true, CancellationToken.None);
    }

    /// <summary>
    /// 接收单条消息（用于认证等握手场景）
    /// </summary>
    private async Task<HubMessage?> ReceiveOneMessageAsync(CancellationToken ct)
    {
        if (_ws?.State != WebSocketState.Open) return null;

        var buffer = new byte[65536];
        var messageBuffer = new MemoryStream();
        WebSocketReceiveResult result;
        do
        {
            result = await _ws.ReceiveAsync(new ArraySegment<byte>(buffer), ct);
            if (result.MessageType == WebSocketMessageType.Close) return null;
            messageBuffer.Write(buffer, 0, result.Count);
        }
        while (!result.EndOfMessage);

        if (result.MessageType != WebSocketMessageType.Text) return null;
        var json = Encoding.UTF8.GetString(messageBuffer.ToArray());
        return HubMessage.FromJson(json);
    }

    /// <summary>
    /// 接收消息循环（带分片处理，使用 ArrayPool 复用缓冲区）
    /// </summary>
    private async Task ReceiveLoopAsync(CancellationToken ct)
    {
        var buffer = ArrayPool<byte>.Shared.Rent(65536);
        var messageBuffer = new MemoryStream();

        try
        {
            while (_ws?.State == WebSocketState.Open && !ct.IsCancellationRequested)
            {
                messageBuffer.SetLength(0); // 复用 MemoryStream
                WebSocketReceiveResult result;
                do
                {
                    result = await _ws.ReceiveAsync(new ArraySegment<byte>(buffer), ct);

                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        return; // 触发重连
                    }

                    messageBuffer.Write(buffer, 0, result.Count);
                }
                while (!result.EndOfMessage);

                if (result.MessageType == WebSocketMessageType.Text)
                {
                    var json = Encoding.UTF8.GetString(messageBuffer.ToArray());
                    var message = HubMessage.FromJson(json);

                    if (message != null)
                    {
                        HandleMessage(message);
                    }
                }
            }
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
            messageBuffer.Dispose();
        }
    }

    /// <summary>
    /// 处理接收到的消息
    /// </summary>
    private void HandleMessage(HubMessage message)
    {
        switch (message.Type)
        {
            case "data":
                var data = message.GetPayload<PlcDataModel>();
                if (data != null) DataReceived?.Invoke(data);
                break;

            case "alarm":
                var alarm = message.GetPayload<AlarmRecord>();
                if (alarm != null) AlarmReceived?.Invoke(alarm);
                break;

            case "lock_result":
                var lockResult = message.GetPayload<LockResultPayload>();
                if (lockResult != null)
                {
                    HasLock = lockResult.Success;
                    CurrentOperator = lockResult.CurrentOperator;
                    LockStatusChanged?.Invoke(lockResult.Success, lockResult.CurrentOperator);
                }
                break;

            case "lock_status":
                var lockStatus = message.GetPayload<LockStatusPayload>();
                if (lockStatus != null)
                {
                    LockStatusChanged?.Invoke(lockStatus.IsLocked, lockStatus.CurrentOperator);
                }
                break;

            case "command_result":
                var cmdResult = message.GetPayload<CommandResultPayload>();
                if (cmdResult != null)
                {
                    CommandResult?.Invoke(cmdResult.Success, cmdResult.Message);
                }
                break;

            case "client_count":
                var count = message.Payload?.GetInt32();
                if (count.HasValue) ClientCountChanged?.Invoke(count.Value);
                break;
        }
    }

    public void Dispose()
    {
        if (_disposed) return;
        _disposed = true;

        Log.Information("HubClient 正在释放... _shouldReconnect={Reconnect}, _connectionTask={HasTask}",
            _shouldReconnect, _connectionTask != null);

        _shouldReconnect = false;
        _cts?.Cancel();

        if (_ws?.State == WebSocketState.Open)
        {
            try
            {
                _ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "关闭", CancellationToken.None).Wait(1000);
            }
            catch { /* 断开清理中，忽略关闭异常 */ }
        }

        _ws?.Dispose();

        try
        {
            var completed = _connectionTask?.Wait(TimeSpan.FromSeconds(2));
            if (completed == true)
                Log.Information("HubClient 连接任务已完成");
            else
                Log.Warning("HubClient 连接任务在2秒内未完成");
        }
        catch (AggregateException ex) when (ex.InnerExceptions.All(e => e is OperationCanceledException))
        {
            // 取消是正常退出路径。
            Log.Information("HubClient 连接任务已取消（正常）");
        }
        _connectionTask = null;
        _cts?.Dispose();

        Log.Information("HubClient 已释放");
    }
}