LightlessClient/LightlessSync/PlayerData/Pairs/PairHandlerRegistry.cs

using LightlessSync.API.Data.Extensions;
using LightlessSync.API.Dto.User;
using Microsoft.Extensions.Logging;

namespace LightlessSync.PlayerData.Pairs;

/// <summary>
///     creates, tracks, and removes pair handlers
/// </summary>
public sealed class PairHandlerRegistry : IDisposable
{
    private readonly object _gate = new();
    private readonly object _pendingGate = new();
    private readonly object _visibilityGate = new();
    private readonly Dictionary<string, PairHandlerEntry> _entriesByIdent = new(StringComparer.Ordinal);
    private readonly Dictionary<string, CancellationTokenSource> _pendingInvisibleEvictions = new(StringComparer.Ordinal);
    private readonly Dictionary<IPairHandlerAdapter, PairHandlerEntry> _entriesByHandler = new(ReferenceEqualityComparer.Instance);

    private readonly IPairHandlerAdapterFactory _handlerFactory;
    private readonly PairManager _pairManager;
    private readonly PairStateCache _pairStateCache;
    private readonly PairPerformanceMetricsCache _pairPerformanceMetricsCache;
    private readonly ILogger<PairHandlerRegistry> _logger;

    private readonly TimeSpan _deletionGracePeriod = TimeSpan.FromMinutes(5);
    private static readonly TimeSpan _handlerReadyTimeout = TimeSpan.FromMinutes(3);
    private const int _handlerReadyPollDelayMs = 500;
    private readonly Dictionary<string, CancellationTokenSource> _pendingCharacterData = new(StringComparer.Ordinal);

    public PairHandlerRegistry(
        IPairHandlerAdapterFactory handlerFactory,
        PairManager pairManager,
        PairStateCache pairStateCache,
        PairPerformanceMetricsCache pairPerformanceMetricsCache,
        ILogger<PairHandlerRegistry> logger)
    {
        _handlerFactory = handlerFactory;
        _pairManager = pairManager;
        _pairStateCache = pairStateCache;
        _pairPerformanceMetricsCache = pairPerformanceMetricsCache;
        _logger = logger;
    }

    public int GetVisibleUsersCount()
    {
        lock (_gate)
        {
            return _entriesByHandler.Keys.Count(handler => handler.IsVisible);
        }
    }

    public bool IsIdentVisible(string ident)
    {
        lock (_gate)
        {
            return _entriesByIdent.TryGetValue(ident, out var entry) && entry.Handler.IsVisible;
        }
    }

    public PairOperationResult<PairUniqueIdentifier> RegisterOnlinePair(PairRegistration registration)
    {
        if (registration.CharacterIdent is null)
        {
            return PairOperationResult<PairUniqueIdentifier>.Fail($"Registration for {registration.PairIdent.UserId} missing ident.");
        }

        IPairHandlerAdapter handler;
        lock (_gate)
        {
            var entry = GetOrCreateEntry(registration.CharacterIdent);
            handler = entry.Handler;
            handler.ScheduledForDeletion = false;
            entry.AddPair(registration.PairIdent);
            if (!handler.Initialized)
            {
                handler.Initialize();
            }
        }

        ApplyPauseStateForHandler(handler);

        if (handler.LastReceivedCharacterData is null)
        {
            var cachedData = _pairStateCache.TryLoad(registration.CharacterIdent);
            if (cachedData is not null)
            {
                handler.LoadCachedCharacterData(cachedData);
            }
        }

        if (handler.LastReceivedCharacterData is not null &&
            (handler.LastAppliedApproximateVRAMBytes < 0 || handler.LastAppliedDataTris < 0 || handler.LastAppliedApproximateEffectiveTris < 0))
        {
            handler.ApplyLastReceivedData(forced: true);
        }

        return PairOperationResult<PairUniqueIdentifier>.Ok(registration.PairIdent);
    }

    public PairOperationResult<PairUniqueIdentifier> DeregisterOfflinePair(PairRegistration registration, bool forceDisposal = false)
    {
        if (registration.CharacterIdent is null)
        {
            return PairOperationResult<PairUniqueIdentifier>.Fail($"Deregister for {registration.PairIdent.UserId} missing ident.");
        }

        IPairHandlerAdapter? handler = null;
        bool shouldScheduleRemoval = false;
        bool shouldDisposeImmediately = false;

        lock (_gate)
        {
            if (!_entriesByIdent.TryGetValue(registration.CharacterIdent, out var entry))
            {
                return PairOperationResult<PairUniqueIdentifier>.Fail($"Ident {registration.CharacterIdent} not registered.");
            }

            handler = entry.Handler;
            entry.RemovePair(registration.PairIdent);
            if (entry.PairCount == 0)
            {
                if (forceDisposal)
                {
                    shouldDisposeImmediately = true;
                }
                else
                {
                    shouldScheduleRemoval = true;
                    handler.ScheduledForDeletion = true;
                }
            }
        }

        if (shouldDisposeImmediately && handler is not null)
        {
            if (TryFinalizeHandlerRemoval(handler))
            {
                handler.Dispose();
            }
        }
        else if (shouldScheduleRemoval && handler is not null)
        {
            _ = RemoveAfterGracePeriodAsync(handler);
        }

        return PairOperationResult<PairUniqueIdentifier>.Ok(registration.PairIdent);
    }

    private PairOperationResult CancelAllInvisibleEvictions()
    {
        List<CancellationTokenSource> snapshot;
        lock (_visibilityGate)
        {
            snapshot = [.. _pendingInvisibleEvictions.Values];
            _pendingInvisibleEvictions.Clear();
        }

        List<string>? errors = null;

        foreach (var cts in snapshot)
        {
            try { cts.Cancel(); }
            catch (Exception ex)
            {
                (errors ??= new List<string>()).Add($"Cancel: {ex.Message}");
            }

            try { cts.Dispose(); }
            catch (Exception ex)
            {
                (errors ??= new List<string>()).Add($"Dispose: {ex.Message}");
            }
        }

        return errors is null
            ? PairOperationResult.Ok()
            : PairOperationResult.Fail($"CancelAllInvisibleEvictions had error(s): {string.Join(" | ", errors)}");
    }

    public PairOperationResult ApplyCharacterData(PairRegistration registration, OnlineUserCharaDataDto dto)
    {
        if (registration.CharacterIdent is null)
        {
            return PairOperationResult.Fail($"Character data received without ident for {registration.PairIdent.UserId}.");
        }

        if (!TryGetHandler(registration.CharacterIdent, out var handler) || handler is null)
        {
            var registerResult = RegisterOnlinePair(registration);
            if (!registerResult.Success)
            {
                return PairOperationResult.Fail(registerResult.Error);
            }

            if (!TryGetHandler(registration.CharacterIdent, out handler) || handler is null)
            {
                QueuePendingCharacterData(registration, dto);
                return PairOperationResult.Ok();
            }
        }

        if (!handler.Initialized)
        {
            handler.Initialize();
            QueuePendingCharacterData(registration, dto);
            return PairOperationResult.Ok();
        }

        handler.ApplyData(dto.CharaData);
        return PairOperationResult.Ok();
    }

    public PairOperationResult ApplyLastReceivedData(PairUniqueIdentifier pairIdent, string ident, bool forced = false)
    {
        if (!TryGetHandler(ident, out var handler) || handler is null)
        {
            return PairOperationResult.Fail($"Cannot reapply data: handler for {pairIdent.UserId} not found.");
        }

        handler.ApplyLastReceivedData(forced);
        return PairOperationResult.Ok();
    }

    public PairOperationResult SetUploading(PairUniqueIdentifier pairIdent, string ident, bool uploading)
    {
        if (!TryGetHandler(ident, out var handler) || handler is null)
        {
            return PairOperationResult.Fail($"Cannot set uploading for {pairIdent.UserId}: handler not found.");
        }

        handler.SetUploading(uploading);
        return PairOperationResult.Ok();
    }

    public PairOperationResult SetPausedState(PairUniqueIdentifier pairIdent, string ident, bool paused)
    {
        if (!TryGetHandler(ident, out var handler) || handler is null)
        {
            return PairOperationResult.Fail($"Cannot update pause state for {pairIdent.UserId}: handler not found.");
        }

        _ = paused; // value reflected in pair manager already
        ApplyPauseStateForHandler(handler);
        return PairOperationResult.Ok();
    }

    public PairOperationResult<IReadOnlyList<(PairUniqueIdentifier Ident, PairConnection Pair)>> GetPairConnections(string ident)
    {
        PairHandlerEntry? entry;
        lock (_gate)
        {
            _entriesByIdent.TryGetValue(ident, out entry);
        }

        if (entry is null)
        {
            return PairOperationResult<IReadOnlyList<(PairUniqueIdentifier Ident, PairConnection Pair)>>.Fail($"No handler registered for {ident}.");
        }

        var list = new List<(PairUniqueIdentifier, PairConnection)>();
        foreach (var pairIdent in entry.SnapshotPairs())
        {
            var result = _pairManager.GetPair(pairIdent.UserId);
            if (result.Success)
            {
                list.Add((pairIdent, result.Value));
            }
        }

        return PairOperationResult<IReadOnlyList<(PairUniqueIdentifier Ident, PairConnection Pair)>>.Ok(list);
    }

    private void ApplyPauseStateForHandler(IPairHandlerAdapter handler)
    {
        var pairs = _pairManager.GetPairsByIdent(handler.Ident);
        bool paused = pairs.Any(p => p.SelfToOtherPermissions.IsPaused() || p.OtherToSelfPermissions.IsPaused());
        handler.SetPaused(paused);
    }

    internal bool TryGetHandler(string ident, out IPairHandlerAdapter? handler)
    {
        lock (_gate)
        {
            var success = _entriesByIdent.TryGetValue(ident, out var entry);
            handler = entry?.Handler;
            return success;
        }
    }

    internal IReadOnlyList<IPairHandlerAdapter> GetHandlerSnapshot()
    {
        lock (_gate)
        {
            return _entriesByHandler.Keys.ToList();
        }
    }

    internal IReadOnlyCollection<PairUniqueIdentifier> GetRegisteredPairs(IPairHandlerAdapter handler)
    {
        lock (_gate)
        {
            if (_entriesByHandler.TryGetValue(handler, out var entry))
            {
                return entry.SnapshotPairs();
            }
        }

        return Array.Empty<PairUniqueIdentifier>();
    }

    internal void ReapplyAll(bool forced = false)
    {
        var handlers = GetHandlerSnapshot();
        foreach (var handler in handlers)
        {
            try
            {
                handler.ApplyLastReceivedData(forced);
            }
            catch (Exception ex)
            {
                if (_logger.IsEnabled(LogLevel.Debug))
                {
                    _logger.LogDebug(ex, "Failed to reapply cached data for {Ident}", handler.Ident);
                }
            }
        }
    }

    internal void ResetAllHandlers()
    {
        List<IPairHandlerAdapter> handlers;
        lock (_gate)
        {
            handlers = _entriesByHandler.Keys.ToList();
            CancelAllInvisibleEvictions();
            _entriesByIdent.Clear();
            _entriesByHandler.Clear();
        }

        CancelAllPendingCharacterData();

        foreach (var handler in handlers)
        {
            try
            {
                handler.Dispose();
            }
            catch (Exception ex)
            {
                if (_logger.IsEnabled(LogLevel.Debug))
                {
                    _logger.LogDebug(ex, "Failed to dispose handler for {Ident}", handler.Ident);
                }
            }
            finally
            {
                _pairPerformanceMetricsCache.Clear(handler.Ident);
            }
        }
    }

    public void Dispose()
    {
        List<IPairHandlerAdapter> handlers;
        lock (_gate)
        {
            handlers = _entriesByHandler.Keys.ToList();
            CancelAllInvisibleEvictions();
            _entriesByIdent.Clear();
            _entriesByHandler.Clear();
        }

        CancelAllPendingCharacterData();

        foreach (var handler in handlers)
        {
            handler.Dispose();
            _pairPerformanceMetricsCache.Clear(handler.Ident);
        }
    }

    private PairHandlerEntry GetOrCreateEntry(string ident)
    {
        if (_entriesByIdent.TryGetValue(ident, out var entry))
        {
            return entry;
        }

        var handler = _handlerFactory.Create(ident);
        entry = new PairHandlerEntry(ident, handler);
        _entriesByIdent[ident] = entry;
        _entriesByHandler[handler] = entry;
        return entry;
    }

    private async Task RemoveAfterGracePeriodAsync(IPairHandlerAdapter handler)
    {
        await Task.Delay(_deletionGracePeriod).ConfigureAwait(false);

        if (TryFinalizeHandlerRemoval(handler))
        {
            handler.Dispose();
        }
    }

    private bool TryFinalizeHandlerRemoval(IPairHandlerAdapter handler)
    {
        string? ident = null;
        lock (_gate)
        {
            if (!_entriesByHandler.TryGetValue(handler, out var entry) || entry.HasPairs)
            {
                handler.ScheduledForDeletion = false;
                return false;
            }

            ident = entry.Ident;
            _entriesByHandler.Remove(handler);
            _entriesByIdent.Remove(entry.Ident);
        }

        if (ident is not null)
        {
            _pairPerformanceMetricsCache.Clear(ident);
            CancelPendingCharacterData(ident);
        }

        return true;
    }

    private void QueuePendingCharacterData(PairRegistration registration, OnlineUserCharaDataDto dto)
    {
        if (registration.CharacterIdent is null) return;

        CancellationTokenSource? previous;
        CancellationTokenSource cts;

        lock (_pendingGate)
        {
            _pendingCharacterData.TryGetValue(registration.CharacterIdent, out previous);
            previous?.Cancel(); 

            cts = new CancellationTokenSource();
            _pendingCharacterData[registration.CharacterIdent] = cts;
        }

        cts.CancelAfter(_handlerReadyTimeout);
        _ = Task.Run(() => WaitThenApplyPendingCharacterDataAsync(registration, dto, cts.Token, cts));
    }

    private void CancelPendingCharacterData(string ident)
    {
        CancellationTokenSource? cts = null;
        lock (_pendingGate)
        {
            if (_pendingCharacterData.TryGetValue(ident, out cts))
                _pendingCharacterData.Remove(ident);
        }

        cts?.Cancel();
    }

    private void CancelAllPendingCharacterData()
    {
        List<CancellationTokenSource>? snapshot = null;
        lock (_pendingGate)
        {
            if (_pendingCharacterData.Count > 0)
            {
                snapshot = [.. _pendingCharacterData.Values];
                _pendingCharacterData.Clear();
            }
        }

        if (snapshot is null) return;
        foreach (var cts in snapshot) cts.Cancel();
    }

    private async Task WaitThenApplyPendingCharacterDataAsync(
        PairRegistration registration,
        OnlineUserCharaDataDto dto,
        CancellationToken token,
        CancellationTokenSource source)
    {
        if (registration.CharacterIdent is null)
        {
            return;
        }

        try
        {
            while (!token.IsCancellationRequested)
            {
                if (TryGetHandler(registration.CharacterIdent, out var handler) && handler is not null && handler.Initialized)
                {
                    handler.ApplyData(dto.CharaData);
                    break;
                }

                await Task.Delay(_handlerReadyPollDelayMs, token).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException)
        {
            // expected
        }
        finally
        {
            lock (_pendingGate)
            {
                if (_pendingCharacterData.TryGetValue(registration.CharacterIdent, out var current) && ReferenceEquals(current, source))
                {
                    _pendingCharacterData.Remove(registration.CharacterIdent);
                }
            }

            source.Dispose();
        }
    }
}