// 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 = new();
            _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(cancellationToken
                    => Utilities.CreateSyntaxContextWithExistingSpeculativeModelAsync(d, self._position, cancellationToken), cacheResult: true));
            }
        }
    }
}