// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.ExceptionServices;
using System.Threading.Tasks;
using Microsoft.Extensions.Internal;

namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
{
    [DebuggerDisplay("{DebuggerToString(),nq}")]
    [DebuggerTypeProxy(typeof(ServiceProviderEngineScopeDebugView))]
    internal sealed class ServiceProviderEngineScope : IServiceScope, IServiceProvider, IKeyedServiceProvider, IAsyncDisposable, IServiceScopeFactory
    {
        // For testing and debugging only.
        // Entries may be null after BeginDispose nulls out duplicate captures of a shared instance.
        internal IList<object?> Disposables => _disposables ?? (IList<object?>)Array.Empty<object?>();

        // When BeginDispose has at most this many captured disposables, it deduplicates them
        // with an inline reference-equality scan rather than allocating a HashSet.
        private const int MaxDisposablesForLinearDedup = 16;

        private bool _disposed;
        private List<object?>? _disposables;

        public ServiceProviderEngineScope(ServiceProvider provider, bool isRootScope)
        {
            ResolvedServices = new Dictionary<ServiceCacheKey, object?>();
            RootProvider = provider;
            IsRootScope = isRootScope;
        }

        internal Dictionary<ServiceCacheKey, object?> ResolvedServices { get; }

        internal bool Disposed => _disposed;

        // This lock protects state on the scope, in particular, for the root scope, it protects
        // the list of disposable entries only, since ResolvedServices are cached on CallSites
        // For other scopes, it protects ResolvedServices and the list of disposables
        internal object Sync => ResolvedServices;

        public bool IsRootScope { get; }

        internal ServiceProvider RootProvider { get; }

        public object? GetService(Type serviceType)
        {
            if (_disposed)
            {
                ThrowHelper.ThrowObjectDisposedException();
            }

            return RootProvider.GetService(ServiceIdentifier.FromServiceType(serviceType), this);
        }

        public object? GetKeyedService(Type serviceType, object? serviceKey)
        {
            if (_disposed)
            {
                ThrowHelper.ThrowObjectDisposedException();
            }

            return RootProvider.GetKeyedService(serviceType, serviceKey, this);
        }

        public object GetRequiredKeyedService(Type serviceType, object? serviceKey)
        {
            if (_disposed)
            {
                ThrowHelper.ThrowObjectDisposedException();
            }

            return RootProvider.GetRequiredKeyedService(serviceType, serviceKey, this);
        }

        public IServiceProvider ServiceProvider => this;

        public IServiceScope CreateScope() => RootProvider.CreateScope();

        [return: NotNullIfNotNull(nameof(service))]
        internal object? CaptureDisposable(object? service)
        {
            if (ReferenceEquals(this, service) || !(service is IDisposable || service is IAsyncDisposable))
            {
                return service;
            }

            bool disposed = false;
            lock (Sync)
            {
                if (_disposed)
                {
                    disposed = true;
                }
                else
                {
                    _disposables ??= new List<object?>();
                    _disposables.Add(service);
                }
            }

            // Don't run customer code under the lock
            if (disposed)
            {
                if (service is IDisposable disposable)
                {
                    disposable.Dispose();
                }
                else
                {
                    // sync over async, for the rare case that an object only implements IAsyncDisposable and may end up starving the thread pool.
                    object? localService = service; // copy to avoid closure on other paths
                    Task.Run(() => ((IAsyncDisposable)localService).DisposeAsync().AsTask()).GetAwaiter().GetResult();
                }

                ThrowHelper.ThrowObjectDisposedException();
            }

            return service;
        }

        public void Dispose()
        {
            List<object?>? toDispose = BeginDispose();
            if (toDispose is null)
            {
                return;
            }

            object? exceptionsCache = null;
            for (var i = toDispose.Count - 1; i >= 0; i--)
            {
                object? disposableEntry = toDispose[i];
                if (disposableEntry is null)
                {
                    continue;
                }

                try
                {
                    if (disposableEntry is IDisposable disposable)
                    {
                        disposable.Dispose();
                    }
                    else
                    {
                        throw new InvalidOperationException(SR.Format(SR.AsyncDisposableServiceDispose, TypeNameHelper.GetTypeDisplayName(disposableEntry)));
                    }
                }
                catch (Exception exception)
                {
                    AddExceptionToCache(ref exceptionsCache, exception);
                }
            }

            CheckExceptionCache(exceptionsCache);
        }

        public ValueTask DisposeAsync()
        {
            List<object?>? toDispose = BeginDispose();
            if (toDispose is null)
            {
                return default;
            }

            object? exceptionsCache = null;
            for (var i = toDispose.Count - 1; i >= 0; i--)
            {
                object? disposable = toDispose[i];
                if (disposable is null)
                {
                    continue;
                }

                try
                {
                    if (disposable is IAsyncDisposable asyncDisposable)
                    {
                        ValueTask vt = asyncDisposable.DisposeAsync();
                        if (!vt.IsCompletedSuccessfully)
                        {
                            return Await(i, vt, toDispose, exceptionsCache);
                        }

                        // If its a IValueTaskSource backed ValueTask,
                        // inform it its result has been read so it can reset
                        vt.GetAwaiter().GetResult();
                    }
                    else
                    {
                        ((IDisposable)disposable).Dispose();
                    }
                }
                catch (Exception exception)
                {
                    AddExceptionToCache(ref exceptionsCache, exception);
                }
            }

            CheckExceptionCache(exceptionsCache);

            return default;

            static async ValueTask Await(int i, ValueTask vt, List<object?> toDispose, object? exceptionsCache)
            {
                try
                {
                    await vt.ConfigureAwait(false);
                }
                catch (Exception exception)
                {
                    AddExceptionToCache(ref exceptionsCache, exception);
                }

                // vt is acting on the disposable at index i,
                // decrement it and move to the next iteration
                i--;

                for (; i >= 0; i--)
                {
                    try
                    {
                        object? disposable = toDispose[i];
                        if (disposable is null)
                        {
                            continue;
                        }

                        if (disposable is IAsyncDisposable asyncDisposable)
                        {
                            await asyncDisposable.DisposeAsync().ConfigureAwait(false);
                        }
                        else
                        {
                            ((IDisposable)disposable).Dispose();
                        }
                    }
                    catch (Exception exception)
                    {
                        AddExceptionToCache(ref exceptionsCache, exception);
                    }
                }

                CheckExceptionCache(exceptionsCache);
            }
        }

        private static void AddExceptionToCache(ref object? exceptionsCache, Exception exception)
        {
            if (exceptionsCache is null)
            {
                exceptionsCache = ExceptionDispatchInfo.Capture(exception);
            }
            else if (exceptionsCache is ExceptionDispatchInfo exceptionInfo)
            {
                exceptionsCache = new List<Exception> { exceptionInfo.SourceException, exception };
            }
            else
            {
                ((List<Exception>)exceptionsCache).Add(exception);
            }
        }

        private static void CheckExceptionCache(object? exceptionsCache)
        {
            if (exceptionsCache is null)
            {
                return;
            }

            if (exceptionsCache is ExceptionDispatchInfo exceptionInfo)
            {
                exceptionInfo.Throw();
            }

            throw new AggregateException((List<Exception>)exceptionsCache);
        }

        private List<object?>? BeginDispose()
        {
            lock (Sync)
            {
                if (_disposed)
                {
                    return null;
                }

                // Track statistics about the scope (number of disposable objects and number of disposed services)
                DependencyInjectionEventSource.Log.ScopeDisposed(RootProvider.GetHashCode(), ResolvedServices.Count, _disposables?.Count ?? 0);

                // We've transitioned to the disposed state, so future calls to
                // CaptureDisposable will immediately dispose the object.
                // No further changes to _state.Disposables, are allowed.
                _disposed = true;

            }

            if (IsRootScope && !RootProvider.IsDisposed())
            {
                // If this ServiceProviderEngineScope instance is a root scope, disposing this instance will need to dispose the RootProvider too.
                // Otherwise the RootProvider will never get disposed and will leak.
                // Note, if the RootProvider get disposed first, it will automatically dispose all attached ServiceProviderEngineScope objects.
                RootProvider.Dispose();
            }

            // ResolvedServices is never cleared for singletons because there might be a compilation running in background
            // trying to get a cached singleton service. If it doesn't find it
            // it will try to create a new one which will result in an ObjectDisposedException.
            if (_disposables is not { Count: > 0 } disposables)
            {
                return null;
            }

            // Lazily deduplicate captured disposables by reference so a shared singleton resolved
            // via multiple factory registrations is disposed exactly once per scope. Later duplicates
            // are nulled in place; the disposal loops walk the list in reverse and skip nulls,
            // which preserves the existing dependent-before-shared ordering.
            DeduplicateDisposables(disposables);

            return disposables;
        }

        private static void DeduplicateDisposables(List<object?> disposables)
        {
            int count = disposables.Count;

            if (count > MaxDisposablesForLinearDedup)
            {
#if NETFRAMEWORK || NETSTANDARD2_0
                var seen = new HashSet<object>(ReferenceEqualityComparer.Instance);
#else
                var seen = new HashSet<object>(count, ReferenceEqualityComparer.Instance);
#endif
                for (int i = 0; i < count; i++)
                {
                    object? entry = disposables[i];
                    if (entry is null)
                    {
                        continue;
                    }

                    if (!seen.Add(entry))
                    {
                        disposables[i] = null;
                    }
                }
            }
            else
            {
                for (int i = 0; i < count; i++)
                {
                    object? entry = disposables[i];
                    if (entry is null)
                    {
                        continue;
                    }

                    for (int j = i + 1; j < count; j++)
                    {
                        if (ReferenceEquals(entry, disposables[j]))
                        {
                            disposables[j] = null;
                        }
                    }
                }
            }
        }

        internal string DebuggerToString()
        {
            string debugText = $"ServiceDescriptors = {RootProvider.CallSiteFactory.Descriptors.Length}";
            if (!IsRootScope)
            {
                debugText += $", IsScope = true";
            }
            if (_disposed)
            {
                debugText += $", Disposed = true";
            }
            return debugText;
        }

        private sealed class ServiceProviderEngineScopeDebugView
        {
            private readonly ServiceProviderEngineScope _serviceProvider;

            public ServiceProviderEngineScopeDebugView(ServiceProviderEngineScope serviceProvider)
            {
                _serviceProvider = serviceProvider;
            }

            public List<ServiceDescriptor> ServiceDescriptors => new List<ServiceDescriptor>(_serviceProvider.RootProvider.CallSiteFactory.Descriptors);
            public List<object> Disposables => FilterDisposables(_serviceProvider.Disposables);
            public bool Disposed => _serviceProvider._disposed;
            public bool IsScope => !_serviceProvider.IsRootScope;

            private static List<object> FilterDisposables(IList<object?> source)
            {
                var result = new List<object>(source.Count);
                for (int i = 0; i < source.Count; i++)
                {
                    if (source[i] is object item)
                    {
                        result.Add(item);
                    }
                }
                return result;
            }
        }
    }
}