// 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.Immutable;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Shared.Extensions.ContextQuery;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.Completion;
 
internal sealed class SharedSyntaxContextsWithSpeculativeModel
{
    private readonly Document _document;
    private readonly int _position;
 
    private readonly ConcurrentDictionary<Document, AsyncLazy<SyntaxContext>> _cache;
    private readonly Lazy<ImmutableArray<DocumentId>> _lazyRelatedDocumentIds;
 
    public SharedSyntaxContextsWithSpeculativeModel(Document document, int position)
    {
        _document = document;
        _position = position;
        _cache = [];
        _lazyRelatedDocumentIds = new(_document.GetLinkedDocumentIds, isThreadSafe: true);
    }
 
    public Task<SyntaxContext> GetSyntaxContextAsync(Document document, CancellationToken cancellationToken)
    {
        if (!_cache.TryGetValue(document, out var lazyContext))
        {
            if (_document.Id != document.Id && !_lazyRelatedDocumentIds.Value.Contains(document.Id))
                throw new ArgumentException("Don't support getting SyntaxContext for document unrelated to the original document");
 
            lazyContext = GetLazySyntaxContextWithSpeculativeModel(document, this);
        }
 
        return lazyContext.GetValueAsync(cancellationToken);
 
        // Extract a local function to avoid creating a closure for code path of cache hit.
        static AsyncLazy<SyntaxContext> GetLazySyntaxContextWithSpeculativeModel(Document document, SharedSyntaxContextsWithSpeculativeModel self)
        {
            return self._cache.GetOrAdd(document, d => AsyncLazy.Create(static (arg, cancellationToken)
                => Utilities.CreateSyntaxContextWithExistingSpeculativeModelAsync(arg.d, arg._position, cancellationToken), (d, self._position)));
        }
    }
}