﻿// 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.Generic;
using System.Collections.Immutable;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Collections;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Microsoft.CodeAnalysis.Shared.TestHooks;
using Microsoft.CodeAnalysis.Shared.Utilities;
using Microsoft.VisualStudio.LanguageServer.Protocol;
using Roslyn.Utilities;

namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SemanticTokens;

/// <summary>
/// Batches requests to refresh the semantic tokens to optomize user experience.
/// </summary>
/// <remarks>This implements <see cref="IOnInitialized"/> to avoid race conditions
/// related to creating the queue on the first request.</remarks>
internal class SemanticTokensRefreshQueue :
    IOnInitialized,
    ILspService,
    IDisposable
{
    /// <summary>
    /// Lock over the mutable state that follows.
    /// </summary>
    private readonly object _gate = new();

    /// <summary>
    /// Mapping from project id to the workqueue for producing the corresponding compilation for it on the OOP server.
    /// </summary>
    private readonly Dictionary<ProjectId, CompilationAvailableEventSource> _projectIdToEventSource = new();

    /// <summary>
    /// Mapping from project id to the project-cone-checksum for it we were at when the project for it had its
    /// compilation produced on the oop server.
    /// </summary>
    private readonly Dictionary<ProjectId, Checksum> _projectIdToLastComputedChecksum = new();

    private readonly LspWorkspaceManager _lspWorkspaceManager;
    private readonly IClientLanguageServerManager _notificationManager;

    private readonly IAsynchronousOperationListener _asyncListener;
    private readonly CancellationTokenSource _disposalTokenSource;

    /// <summary>
    /// Debouncing queue so that we don't attempt to issue a semantic tokens refresh notification too often.
    /// 
    /// Null when the client does not support sending refresh notifications.
    /// </summary>
    private AsyncBatchingWorkQueue<Uri?>? _semanticTokenRefreshQueue;

    private readonly LspWorkspaceRegistrationService _lspWorkspaceRegistrationService;

    public SemanticTokensRefreshQueue(
        IAsynchronousOperationListenerProvider asynchronousOperationListenerProvider,
        LspWorkspaceRegistrationService lspWorkspaceRegistrationService,
        LspWorkspaceManager lspWorkspaceManager,
        IClientLanguageServerManager notificationManager)
    {
        _asyncListener = asynchronousOperationListenerProvider.GetListener(FeatureAttribute.Classification);

        _lspWorkspaceRegistrationService = lspWorkspaceRegistrationService;
        _disposalTokenSource = new();

        _lspWorkspaceManager = lspWorkspaceManager;
        _notificationManager = notificationManager;
    }

    public Task OnInitializedAsync(ClientCapabilities clientCapabilities, CancellationToken _)
    {
        if (_semanticTokenRefreshQueue is null && clientCapabilities.Workspace?.SemanticTokens?.RefreshSupport is true)
        {
            // Only send a refresh notification to the client every 2s (if needed) in order to avoid
            // sending too many notifications at once.  This ensures we batch up workspace notifications,
            // but also means we send soon enough after a compilation-computation to not make the user wait
            // an enormous amount of time.
            _semanticTokenRefreshQueue = new AsyncBatchingWorkQueue<Uri?>(
                delay: TimeSpan.FromMilliseconds(2000),
                processBatchAsync: (documentUris, cancellationToken)
                    => FilterLspTrackedDocumentsAsync(_lspWorkspaceManager, _notificationManager, documentUris, cancellationToken),
                equalityComparer: EqualityComparer<Uri?>.Default,
                asyncListener: _asyncListener,
                _disposalTokenSource.Token);

            _lspWorkspaceRegistrationService.LspSolutionChanged += OnLspSolutionChanged;
        }

        return Task.CompletedTask;
    }

    public async Task TryEnqueueRefreshComputationAsync(Project project, CancellationToken cancellationToken)
    {
        if (_semanticTokenRefreshQueue is not null)
        {
            // Determine the checksum for this project cone.  Note: this should be fast in practice because this is
            // the same project-cone-checksum we used to even call into OOP above when we computed semantic tokens.
            var projectChecksum = await project.Solution.State.GetChecksumAsync(project.Id, cancellationToken).ConfigureAwait(false);

            lock (_gate)
            {
                // If this checksum is the same as the last computed result, no need to continue, we would not produce a
                // different compilation.
                if (ChecksumIsUnchanged_NoLock(project, projectChecksum))
                    return;

                if (!_projectIdToEventSource.TryGetValue(project.Id, out var eventSource))
                {
                    eventSource = new CompilationAvailableEventSource(_asyncListener);
                    _projectIdToEventSource.Add(project.Id, eventSource);
                }

                eventSource.EnsureCompilationAvailability(project, () => OnCompilationAvailable(project, projectChecksum));
            }
        }
    }

    private static ValueTask FilterLspTrackedDocumentsAsync(
        LspWorkspaceManager lspWorkspaceManager,
        IClientLanguageServerManager notificationManager,
        ImmutableSegmentedList<Uri?> documentUris,
        CancellationToken cancellationToken)
    {
        var trackedDocuments = lspWorkspaceManager.GetTrackedLspText();
        foreach (var documentUri in documentUris)
        {
            if (documentUri is null || !trackedDocuments.ContainsKey(documentUri))
            {
                return notificationManager.SendRequestAsync(Methods.WorkspaceSemanticTokensRefreshName, cancellationToken);
            }
        }

        // LSP is already tracking all changed documents so we don't need to send a refresh request.
        return ValueTaskFactory.CompletedTask;
    }

    private void OnLspSolutionChanged(object? sender, WorkspaceChangeEventArgs e)
    {
        if (e.DocumentId is not null && e.Kind is WorkspaceChangeKind.DocumentChanged)
        {
            var document = e.NewSolution.GetRequiredDocument(e.DocumentId);
            var documentUri = document.GetURI();

            // We enqueue the URI since there's a chance the client is already tracking the
            // document, in which case we don't need to send a refresh notification.
            // We perform the actual check when processing the batch to ensure we have the
            // most up-to-date list of tracked documents.
            EnqueueSemanticTokenRefreshNotification(documentUri);
        }
        else
        {
            EnqueueSemanticTokenRefreshNotification(documentUri: null);
        }
    }

    private void EnqueueSemanticTokenRefreshNotification(Uri? documentUri)
    {
        // We should have only gotten here if semantic tokens refresh is supported and initialized.
        Contract.ThrowIfNull(_semanticTokenRefreshQueue);
        _semanticTokenRefreshQueue.AddWork(documentUri);
    }

    private void OnCompilationAvailable(Project project, Checksum projectChecksum)
    {
        lock (_gate)
        {
            // Paranoia: It's technically possible (though unlikely) for two calls to compute the compilation for
            // the same project to come back and call into this.  This is because the
            // CompilationAvailableEventSource uses cooperative cancellation to cancel the in-flight request before
            // issuing the new one.  There is no requirement though that the inflight request actually stop (or run
            // to completion) prior to the next request running and completing.  In practice this should not happen
            // as cancellation is checked fairly regularly.  However, if it does, check and do not bother to issue a
            // refresh in this case.
            if (ChecksumIsUnchanged_NoLock(project, projectChecksum))
                return;

            // keep track of this checksum.  That way we don't get into a loop where we send a refresh notification,
            // then we get called back into, causing us to compute the compilation, causing us to send the refresh
            // notification, etc. etc.
            _projectIdToLastComputedChecksum[project.Id] = projectChecksum;
        }

        EnqueueSemanticTokenRefreshNotification(documentUri: null);
    }

    private bool ChecksumIsUnchanged_NoLock(Project project, Checksum projectChecksum)
        => _projectIdToLastComputedChecksum.TryGetValue(project.Id, out var lastChecksum) && lastChecksum == projectChecksum;

    public void Dispose()
    {
        ImmutableArray<CompilationAvailableEventSource> eventSources;
        lock (_gate)
        {
            eventSources = _projectIdToEventSource.Values.ToImmutableArray();
            _projectIdToEventSource.Clear();

            _lspWorkspaceRegistrationService.LspSolutionChanged -= OnLspSolutionChanged;
        }

        foreach (var eventSource in eventSources)
            eventSource.Dispose();

        _disposalTokenSource.Cancel();
        _disposalTokenSource.Dispose();
    }
}
