// 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.Collections.Generic;
using System.Composition;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Host;
using Microsoft.CodeAnalysis.Host.Mef;
using Microsoft.CodeAnalysis.Test.Utilities;
using Microsoft.CodeAnalysis.UnitTests.Remote;
using Microsoft.VisualStudio.Threading;
using Roslyn.Test.Utilities;
 
namespace Microsoft.CodeAnalysis.Remote.Testing;
 
#pragma warning disable CA1416 // Validate platform compatibility
internal sealed class InProcRemoteHostClientProvider : IRemoteHostClientProvider, IDisposable
{
    [ExportWorkspaceServiceFactory(typeof(IRemoteHostClientProvider), ServiceLayer.Test), Shared, PartNotDiscoverable]
    internal sealed class Factory : IWorkspaceServiceFactory
    {
        private readonly RemoteServiceCallbackDispatcherRegistry _callbackDispatchers;
 
        [ImportingConstructor]
        [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
        public Factory([ImportMany] IEnumerable<Lazy<IRemoteServiceCallbackDispatcher, RemoteServiceCallbackDispatcherRegistry.ExportMetadata>> callbackDispatchers)
            => _callbackDispatchers = new RemoteServiceCallbackDispatcherRegistry(callbackDispatchers);
 
        public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
            => new InProcRemoteHostClientProvider(workspaceServices.SolutionServices, _callbackDispatchers);
    }
 
    private sealed class WorkspaceManager : RemoteWorkspaceManager
    {
        public WorkspaceManager(
            Func<RemoteWorkspace, SolutionAssetCache> createAssetStorage,
            ConcurrentDictionary<Guid, TestGeneratorReference> sharedTestGeneratorReferences,
            Type[]? additionalRemoteParts,
            Type[]? excludedRemoteParts)
            : base(createAssetStorage, CreateRemoteWorkspace(sharedTestGeneratorReferences, additionalRemoteParts, excludedRemoteParts))
        {
        }
    }
 
    private static RemoteWorkspace CreateRemoteWorkspace(
        ConcurrentDictionary<Guid, TestGeneratorReference> sharedTestGeneratorReferences,
        Type[]? additionalRemoteParts,
        Type[]? excludedRemoteParts)
    {
        var hostServices = FeaturesTestCompositions.RemoteHost.AddParts(additionalRemoteParts).AddExcludedPartTypes(excludedRemoteParts).GetHostServices();
 
        // We want to allow references to source generators to be shared between the "in proc" and "remote" workspaces and
        // MEF compositions, so tell the serializer service to use the same map for this "remote" workspace as the in-proc one.
        ((IMefHostExportProvider)hostServices).GetExportedValue<TestSerializerService.Factory>().SharedTestGeneratorReferences = sharedTestGeneratorReferences;
        return new RemoteWorkspace(hostServices);
    }
 
    private readonly SolutionServices _services;
    private readonly Lazy<WorkspaceManager> _lazyManager;
    private readonly Lazy<RemoteHostClient> _lazyClient;
    private readonly TaskCompletionSource<bool> _clientCreationSource = new();
 
    public Type[]? AdditionalRemoteParts { get; set; }
    public Type[]? ExcludedRemoteParts { get; set; }
    public TraceListener? TraceListener { get; set; }
 
    public InProcRemoteHostClientProvider(SolutionServices services, RemoteServiceCallbackDispatcherRegistry callbackDispatchers)
    {
        _services = services;
 
        var testSerializerServiceFactory = services.ExportProvider.GetExportedValue<TestSerializerService.Factory>();
 
        _lazyManager = new Lazy<WorkspaceManager>(
            () => new WorkspaceManager(
                _ => new SolutionAssetCache(),
                testSerializerServiceFactory.SharedTestGeneratorReferences,
                AdditionalRemoteParts,
                ExcludedRemoteParts));
        _lazyClient = new Lazy<RemoteHostClient>(
            () =>
            {
                try
                {
                    return InProcRemoteHostClient.Create(
                        _services,
                        callbackDispatchers,
                        TraceListener,
                        new RemoteHostTestData(_lazyManager.Value, isInProc: true));
                }
                finally
                {
                    _clientCreationSource.SetResult(true);
                }
            });
    }
 
    public void Dispose()
    {
        // Dispose the remote workspace when the owning host workspace is disposed
        if (_lazyManager.IsValueCreated)
        {
            var manager = _lazyManager.Value;
            manager.GetWorkspace().Dispose();
        }
    }
 
    public Task<RemoteHostClient?> TryGetRemoteHostClientAsync(CancellationToken cancellationToken)
        => Task.FromResult<RemoteHostClient?>(_lazyClient.Value);
 
    public Task WaitForClientCreationAsync(CancellationToken cancellationToken)
        => _clientCreationSource.Task.WithCancellation(cancellationToken);
}
#pragma warning restore CA1416 // Validate platform compatibility