Windows-Gotify-Client/Services/GotifyWebSocketService.cs

using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using WindowsGotify.Models;

namespace WindowsGotify.Services;

public sealed class GotifyWebSocketService
{
    private readonly AppConfig _config;
    private readonly Logger _logger;
    private readonly StateService _stateService;
    private readonly GotifyApiService _apiService;
    private readonly NotificationService _notificationService;
    private readonly Action<string> _statusChanged;

    private readonly CancellationTokenSource _lifetimeCts = new();
    private readonly object _socketLock = new();
    private ClientWebSocket? _socket;
    private bool _running;
    private bool _forceReconnect;

    /// <summary>
    /// 收到一条新消息(已去重、未越界)时触发,UI 用来追加到历史列表。
    /// </summary>
    public event Action<GotifyMessage>? MessageReceived;

    public GotifyWebSocketService(
        AppConfig config,
        Logger logger,
        StateService stateService,
        GotifyApiService apiService,
        NotificationService notificationService,
        Action<string> statusChanged)
    {
        _config = config;
        _logger = logger;
        _stateService = stateService;
        _apiService = apiService;
        _notificationService = notificationService;
        _statusChanged = statusChanged;
    }

    public async Task StartAsync()
    {
        if (_running)
            return;

        _running = true;
        await _stateService.LoadAsync();
        _ = Task.Run(RunLoopAsync);
    }

    public void RequestReconnect()
    {
        _forceReconnect = true;
        lock (_socketLock)
        {
            try
            {
                _socket?.Abort();
            }
            catch
            {
                // ignored
            }
        }
    }

    public void Stop()
    {
        _running = false;
        _lifetimeCts.Cancel();

        lock (_socketLock)
        {
            try
            {
                _socket?.Abort();
                _socket?.Dispose();
            }
            catch
            {
                // ignored
            }

            _socket = null;
        }
    }

    private async Task RunLoopAsync()
    {
        var delay = Math.Max(1, _config.Reconnect.InitialDelaySeconds);
        var maxDelay = Math.Max(delay, _config.Reconnect.MaxDelaySeconds);
        var firstConnection = true;

        while (_running && !_lifetimeCts.IsCancellationRequested)
        {
            try
            {
                _statusChanged("连接中...");

                if (firstConnection && _config.History.Enabled)
                {
                    await SyncHistoryAsync(_lifetimeCts.Token);
                    firstConnection = false;
                }

                await ConnectAndReceiveAsync(_lifetimeCts.Token);

                delay = Math.Max(1, _config.Reconnect.InitialDelaySeconds);
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (Exception ex)
            {
                _logger.Error("Gotify WebSocket loop failed.", ex);
                _statusChanged("断开");
            }

            if (!_running || _lifetimeCts.IsCancellationRequested)
                break;

            _forceReconnect = false;
            _logger.Info($"Reconnect after {delay}s.");
            _statusChanged($"{delay}s 后重连");

            try
            {
                await Task.Delay(TimeSpan.FromSeconds(delay), _lifetimeCts.Token);
            }
            catch (OperationCanceledException)
            {
                break;
            }

            delay = Math.Min(maxDelay, Math.Max(delay + 1, delay * 2));
        }
    }

    private async Task SyncHistoryAsync(CancellationToken ct)
    {
        try
        {
            var lastId = _stateService.LastMessageId;
            _logger.Info($"History sync. lastMessageId={lastId}");

            var messages = await _apiService.GetMessagesSinceAsync(
                lastId,
                _config.History.StartupSyncCount,
                ct);

            foreach (var message in messages.OrderBy(x => x.Id))
            {
                await ProcessMessageAsync(message, showNotification: message.Id > lastId, ct);
            }
        }
        catch (Exception ex)
        {
            _logger.Error("History sync failed.", ex);
        }
    }

    private async Task ConnectAndReceiveAsync(CancellationToken ct)
    {
        using var socket = new ClientWebSocket();
        var wsUri = BuildWebSocketUri();

        lock (_socketLock)
        {
            _socket = socket;
        }

        _logger.Info($"Connecting to {HideToken(wsUri)}");

        await socket.ConnectAsync(wsUri, ct);
        _logger.Info("WebSocket connected.");
        _statusChanged("已连接");

        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        var buffer = new byte[64 * 1024];

        while (socket.State == WebSocketState.Open && !ct.IsCancellationRequested)
        {
            using var ms = new MemoryStream();
            WebSocketReceiveResult result;

            do
            {
                result = await socket.ReceiveAsync(buffer, linkedCts.Token);

                if (result.MessageType == WebSocketMessageType.Close)
                {
                    _logger.Warn("Gotify requested WebSocket close.");
                    return;
                }

                if (result.MessageType == WebSocketMessageType.Text && result.Count > 0)
                {
                    ms.Write(buffer, 0, result.Count);
                }
            }
            while (!result.EndOfMessage);

            if (ms.Length == 0)
                continue;

            var json = Encoding.UTF8.GetString(ms.ToArray());
            await HandleMessageJsonAsync(json, ct);
        }
    }

    private async Task HandleMessageJsonAsync(string json, CancellationToken ct)
    {
        try
        {
            var message = JsonSerializer.Deserialize<GotifyMessage>(json);
            if (message is null)
                return;

            await ProcessMessageAsync(message, showNotification: true, ct);
        }
        catch (JsonException ex)
        {
            _logger.Error("Invalid Gotify message JSON.", ex);
        }
    }

    private async Task ProcessMessageAsync(
        GotifyMessage message,
        bool showNotification,
        CancellationToken ct)
    {
        if (message.Id <= 0)
            return;

        var lastId = _stateService.LastMessageId;
        if (message.Id <= lastId)
        {
            _logger.Info($"Duplicate/old message ignored. id={message.Id}, last={lastId}");
            return;
        }

        if (showNotification)
        {
            await _notificationService.ShowAsync(message);
        }

        await _stateService.SetLastMessageIdAsync(message.Id);
        _logger.Info($"Message processed. id={message.Id}");

        // 触发 UI 订阅者(必须在 await SetLastMessageIdAsync 之后,确保下游不会拿到陈旧 lastId)
        try
        {
            MessageReceived?.Invoke(message);
        }
        catch (Exception ex)
        {
            _logger.Error("MessageReceived subscriber threw.", ex);
        }
    }

    private Uri BuildWebSocketUri()
    {
        var baseUri = new Uri(_config.Gotify.Url.TrimEnd('/') + "/");
        var builder = new UriBuilder(baseUri)
        {
            Scheme = baseUri.Scheme.Equals("https", StringComparison.OrdinalIgnoreCase)
                ? "wss"
                : "ws",
            Path = baseUri.AbsolutePath.TrimEnd('/') + "/stream",
            Query = "token=" + Uri.EscapeDataString(_config.Gotify.ClientToken)
        };

        return builder.Uri;
    }

    private static string HideToken(Uri uri)
    {
        var raw = uri.ToString();
        var marker = "token=";
        var index = raw.IndexOf(marker, StringComparison.OrdinalIgnoreCase);
        return index < 0 ? raw : raw[..(index + marker.Length)] + "***";
    }
}