// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
using System;
using System.Collections.Concurrent;
using System.Threading;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.Remote;
 
internal interface IRemoteServiceCallbackDispatcher
{
    RemoteServiceCallbackDispatcher.Handle CreateHandle(object? instance);
}
 
internal class RemoteServiceCallbackDispatcher : IRemoteServiceCallbackDispatcher
{
    internal readonly struct Handle(ConcurrentDictionary<RemoteServiceCallbackId, object> callbackInstances, RemoteServiceCallbackId callbackId) : IDisposable
    {
        public readonly RemoteServiceCallbackId Id = callbackId;
 
        public void Dispose()
        {
            Contract.ThrowIfTrue(callbackInstances?.TryRemove(Id, out _) == false);
        }
    }
 
    private int _callbackId = 1;
    private readonly ConcurrentDictionary<RemoteServiceCallbackId, object> _callbackInstances = new(concurrencyLevel: 2, capacity: 10);
 
    public Handle CreateHandle(object? instance)
    {
        if (instance is null)
        {
            return default;
        }
 
        var callbackId = new RemoteServiceCallbackId(Interlocked.Increment(ref _callbackId));
        var handle = new Handle(_callbackInstances, callbackId);
        _callbackInstances.Add(callbackId, instance);
        return handle;
    }
 
    public object GetCallback(RemoteServiceCallbackId callbackId)
    {
        Contract.ThrowIfFalse(_callbackInstances.TryGetValue(callbackId, out var instance));
        return instance;
    }
}