PlotDirector/PlotLine/Services/WordCompanionPresenceService.cs

using System.Collections.Concurrent;
using PlotLine.Models;

namespace PlotLine.Services;

public interface IWordCompanionPresenceService
{
    Task<WordCompanionPresenceStatus> RegisterAsync(int userId, string connectionId, WordCompanionPresenceRegistration registration);
    Task<WordCompanionPresenceStatus> HeartbeatAsync(int userId, string connectionId, WordCompanionPresenceHeartbeat heartbeat);
    Task<WordCompanionPresenceStatus?> DisconnectAsync(string connectionId);
    Task<WordCompanionPresenceStatus> GetStatusAsync(int userId);
    Task<bool> IsConnectedAsync(int userId);
    Task<IReadOnlyList<WordCompanionPresenceStatus>> MarkStaleOfflineAsync();
    Task<string?> GetCompanionConnectionIdAsync(int userId);
}

public sealed class WordCompanionPresenceService : IWordCompanionPresenceService
{
    private static readonly TimeSpan HeartbeatTimeout = TimeSpan.FromSeconds(55);
    private readonly ConcurrentDictionary<int, PresenceRecord> recordsByUser = new();
    private readonly ConcurrentDictionary<string, int> usersByConnection = new(StringComparer.Ordinal);

    public Task<WordCompanionPresenceStatus> RegisterAsync(int userId, string connectionId, WordCompanionPresenceRegistration registration)
    {
        var now = DateTime.UtcNow;
        var record = new PresenceRecord
        {
            UserID = userId,
            ConnectionID = connectionId,
            Status = WordCompanionPresenceStatuses.Connected,
            CompanionVersion = Clean(registration.CompanionVersion),
            MachineName = Clean(registration.MachineName),
            DocumentOpen = registration.DocumentOpen,
            CurrentDocumentName = Clean(registration.CurrentDocumentName),
            LinkedProjectID = Positive(registration.LinkedProjectID),
            LinkedBookID = Positive(registration.LinkedBookID),
            ConnectedUtc = now,
            LastHeartbeatUtc = now
        };

        recordsByUser.AddOrUpdate(userId, record, (_, _) => record);
        usersByConnection[connectionId] = userId;
        return Task.FromResult(ToStatus(record, now));
    }

    public Task<WordCompanionPresenceStatus> HeartbeatAsync(int userId, string connectionId, WordCompanionPresenceHeartbeat heartbeat)
    {
        var now = DateTime.UtcNow;
        var record = recordsByUser.AddOrUpdate(
            userId,
            _ => new PresenceRecord
            {
                UserID = userId,
                ConnectionID = connectionId,
                ConnectedUtc = now
            },
            (_, existing) =>
            {
                existing.ConnectionID = connectionId;
                return existing;
            });

        record.Status = WordCompanionPresenceStatuses.Connected;
        record.CompanionVersion = Clean(heartbeat.CompanionVersion) ?? record.CompanionVersion;
        record.MachineName = Clean(heartbeat.MachineName) ?? record.MachineName;
        record.DocumentOpen = heartbeat.DocumentOpen;
        record.CurrentDocumentName = Clean(heartbeat.CurrentDocumentName);
        record.LinkedProjectID = Positive(heartbeat.LinkedProjectID);
        record.LinkedBookID = Positive(heartbeat.LinkedBookID);
        record.LastHeartbeatUtc = now;

        usersByConnection[connectionId] = userId;
        return Task.FromResult(ToStatus(record, now));
    }

    public Task<WordCompanionPresenceStatus?> DisconnectAsync(string connectionId)
    {
        if (!usersByConnection.TryRemove(connectionId, out var userId)
            || !recordsByUser.TryGetValue(userId, out var record)
            || !string.Equals(record.ConnectionID, connectionId, StringComparison.Ordinal))
        {
            return Task.FromResult<WordCompanionPresenceStatus?>(null);
        }

        record.Status = WordCompanionPresenceStatuses.Offline;
        record.ConnectionID = string.Empty;
        return Task.FromResult<WordCompanionPresenceStatus?>(ToStatus(record, DateTime.UtcNow));
    }

    public Task<WordCompanionPresenceStatus> GetStatusAsync(int userId)
    {
        if (!recordsByUser.TryGetValue(userId, out var record))
        {
            return Task.FromResult(new WordCompanionPresenceStatus { UserID = userId });
        }

        return Task.FromResult(ToStatus(record, DateTime.UtcNow));
    }

    public async Task<bool> IsConnectedAsync(int userId)
        => (await GetStatusAsync(userId)).IsConnected;

    public Task<string?> GetCompanionConnectionIdAsync(int userId)
    {
        if (!recordsByUser.TryGetValue(userId, out var record))
        {
            return Task.FromResult<string?>(null);
        }

        var status = ToStatus(record, DateTime.UtcNow);
        return Task.FromResult(status.IsConnected && !string.IsNullOrWhiteSpace(record.ConnectionID)
            ? record.ConnectionID
            : null);
    }

    public Task<IReadOnlyList<WordCompanionPresenceStatus>> MarkStaleOfflineAsync()
    {
        var now = DateTime.UtcNow;
        var changed = new List<WordCompanionPresenceStatus>();
        foreach (var record in recordsByUser.Values)
        {
            if (!string.Equals(record.Status, WordCompanionPresenceStatuses.Connected, StringComparison.Ordinal)
                || !record.LastHeartbeatUtc.HasValue
                || now - record.LastHeartbeatUtc.Value <= HeartbeatTimeout)
            {
                continue;
            }

            record.Status = WordCompanionPresenceStatuses.Offline;
            if (!string.IsNullOrWhiteSpace(record.ConnectionID))
            {
                usersByConnection.TryRemove(record.ConnectionID, out _);
                record.ConnectionID = string.Empty;
            }

            changed.Add(ToStatus(record, now));
        }

        return Task.FromResult<IReadOnlyList<WordCompanionPresenceStatus>>(changed);
    }

    private static WordCompanionPresenceStatus ToStatus(PresenceRecord record, DateTime now)
    {
        var connected = string.Equals(record.Status, WordCompanionPresenceStatuses.Connected, StringComparison.Ordinal)
            && record.LastHeartbeatUtc.HasValue
            && now - record.LastHeartbeatUtc.Value <= HeartbeatTimeout;

        return new WordCompanionPresenceStatus
        {
            UserID = record.UserID,
            IsConnected = connected,
            Status = connected ? WordCompanionPresenceStatuses.Connected : WordCompanionPresenceStatuses.Offline,
            CompanionVersion = record.CompanionVersion,
            MachineName = record.MachineName,
            DocumentOpen = record.DocumentOpen,
            CurrentDocumentName = record.CurrentDocumentName,
            LinkedProjectID = record.LinkedProjectID,
            LinkedBookID = record.LinkedBookID,
            ConnectedUtc = record.ConnectedUtc,
            LastHeartbeatUtc = record.LastHeartbeatUtc
        };
    }

    private static string? Clean(string? value)
    {
        var cleaned = (value ?? string.Empty).Trim();
        return string.IsNullOrWhiteSpace(cleaned) ? null : cleaned;
    }

    private static int? Positive(int? value)
        => value is > 0 ? value.Value : null;

    private sealed class PresenceRecord
    {
        public int UserID { get; init; }
        public string ConnectionID { get; set; } = string.Empty;
        public string Status { get; set; } = WordCompanionPresenceStatuses.Offline;
        public string? CompanionVersion { get; set; }
        public string? MachineName { get; set; }
        public bool DocumentOpen { get; set; }
        public string? CurrentDocumentName { get; set; }
        public int? LinkedProjectID { get; set; }
        public int? LinkedBookID { get; set; }
        public DateTime? ConnectedUtc { get; set; }
        public DateTime? LastHeartbeatUtc { get; set; }
    }
}