using DodoSSH.Api.Setup;
using Microsoft.Extensions.Options;

namespace DodoSSH.Api.Features.Relay;

/// <summary>
/// Counts the relay sessions open on this node, and refuses the ones over ADR 0004's limits.
/// </summary>
/// <remarks>
/// <para>
/// Per node, and so per user <em>per node</em>: a user behind a round-robin balancer with three nodes
/// could hold three times the per-user cap. That is accepted rather than fixed, because fixing it means
/// shared state on the hot path of every connect, and the cap exists to stop one account exhausting a
/// node — which a per-node count does exactly.
/// </para>
/// <para>
/// A singleton, because a session outlives the request that opened it. It holds counts and nothing else;
/// the sockets belong to the handlers that opened them.
/// </para>
/// </remarks>
internal sealed class RelaySessionRegistry(IOptions<RelayOptions> options)
{
    private readonly Lock gate = new();
    private readonly Dictionary<Guid, int> perUser = [];
    private int total;

    /// <summary>How many sessions are open on this node.</summary>
    internal int Count
    {
        get
        {
            lock (gate)
            {
                return total;
            }
        }
    }

    /// <summary>Whether a session for this user would be admitted now.</summary>
    /// <remarks>
    /// Advisory, for the ticket endpoint: it lets a caller already at the cap hear so as an HTTP status
    /// before it has spent a ticket, but only <see cref="TryAdmit"/> decides, since the answer here can
    /// be stale by the time the socket arrives — possibly at another node.
    /// </remarks>
    internal bool HasRoomFor(Guid userId)
    {
        var settings = options.Value;

        lock (gate)
        {
            return total < settings.MaxConcurrentSessionsTotal
                && perUser.GetValueOrDefault(userId) < settings.MaxConcurrentSessionsPerUser;
        }
    }

    /// <summary>Admits a session, or answers null when a limit is met.</summary>
    /// <returns>A lease that gives the place back when disposed.</returns>
    internal Lease? TryAdmit(Guid userId)
    {
        var settings = options.Value;

        lock (gate)
        {
            var held = perUser.GetValueOrDefault(userId);

            if (total >= settings.MaxConcurrentSessionsTotal || held >= settings.MaxConcurrentSessionsPerUser)
            {
                return null;
            }

            perUser[userId] = held + 1;
            total++;

            return new Lease(this, userId);
        }
    }

    private void Release(Guid userId)
    {
        lock (gate)
        {
            total--;

            // Removed at zero rather than left behind, so the dictionary is the size of the users with a
            // session open and not of every user who ever had one.
            if (perUser[userId] <= 1)
            {
                perUser.Remove(userId);
            }
            else
            {
                perUser[userId]--;
            }
        }
    }

    /// <summary>One admitted session's place. Idempotent to dispose.</summary>
    internal sealed class Lease(RelaySessionRegistry registry, Guid userId) : IDisposable
    {
        private int released;

        /// <inheritdoc />
        public void Dispose()
        {
            if (Interlocked.Exchange(ref released, 1) == 0)
            {
                registry.Release(userId);
            }
        }
    }
}