// 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.
 
#nullable disable
 
using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
 
namespace Microsoft.VisualStudio.LanguageServices.Implementation.Interop
{
    internal readonly struct WeakComHandle<THandle, TObject>
        where THandle : class
        where TObject : class, THandle
    {
        // NOTE: The logic here is a little bit tricky.  We can't just keep a WeakReference to
        // something like a ComHandle, since it's not something that our clients keep alive.
        // instead, we keep a weak reference to the inner managed object, which we know will
        // always be alive if the outer aggregate is alive.  We can't just keep a WeakReference
        // to the RCW for the outer object either, since in cases where we have a DCOM or native
        // client, the RCW will be cleaned up, even though there is still a native reference
        // to the underlying native outer object.
        //
        // Instead we make use of an implementation detail of the way the CLR's COM aggregation 
        // works.  Namely, if all references to the aggregated object are released, the CLR 
        // responds to QI's for IUnknown with a different object.  So, we store the original
        // value, when we know that we have a client, and then we use that to compare to see
        // if we still have a client alive.
        //
        // NOTE: This is _NOT_ AddRef'd.  We use it just to store the integer value of the
        // IUnknown for comparison purposes.
        private readonly WeakReference _managedObjectWeakReference;
        private readonly IntPtr _pUnkOfInnerUnknownWhenAlive;
 
        public WeakComHandle(THandle comAggregateObject)
        {
            if (comAggregateObject == null)
            {
                _managedObjectWeakReference = null;
                _pUnkOfInnerUnknownWhenAlive = IntPtr.Zero;
            }
 
            var pUnk = IntPtr.Zero;
            var managedObject = ComAggregate.GetManagedObject<TObject>(comAggregateObject);
 
            try
            {
                pUnk = Marshal.GetIUnknownForObject(managedObject);
                _pUnkOfInnerUnknownWhenAlive = pUnk;
            }
            finally
            {
                if (pUnk != IntPtr.Zero)
                {
                    Marshal.Release(pUnk);
                }
            }
 
            _managedObjectWeakReference = new WeakReference(managedObject);
        }
 
        public WeakComHandle(ComHandle<THandle, TObject> handle)
        {
            var pUnk = IntPtr.Zero;
            try
            {
                pUnk = Marshal.GetIUnknownForObject(handle.Object);
                _pUnkOfInnerUnknownWhenAlive = pUnk;
            }
            finally
            {
                if (pUnk != IntPtr.Zero)
                {
                    Marshal.Release(pUnk);
                }
            }
 
            _managedObjectWeakReference = new WeakReference(handle.Object);
        }
 
        public THandle ComAggregateObject
        {
            get
            {
                // This is pretty fragile code, watch carefully for race conditions!
                var pUnk = IntPtr.Zero;
                try
                {
                    if (_managedObjectWeakReference == null)
                    {
                        return null;
                    }
 
                    // Copy target locally to make sure other thread won't delete it before we use it
                    var target = _managedObjectWeakReference.Target;
                    if (target == null)
                    {
                        return null;
                    }
 
                    pUnk = Marshal.GetIUnknownForObject(target);
                    if (pUnk == _pUnkOfInnerUnknownWhenAlive)
                    {
                        // QueryInterface on COM aggregate might fail during shutdown, so we 
                        // defensively use "as" instead of casting (see Dev10 816848).
                        return Marshal.GetObjectForIUnknown(pUnk) as THandle;
                    }
                    else
                    {
                        return null;
                    }
                }
                finally
                {
                    if (pUnk != IntPtr.Zero)
                    {
                        Marshal.Release(pUnk);
                    }
                }
            }
        }
 
        internal bool TryGetManagedObjectWithoutCaringWhetherNativeObjectIsAlive(out TObject managedObject)
        {
            // NOTE: Only use this method if you do NOT care whether the native ComAggregate
            // object has already been released.
            if (_managedObjectWeakReference == null)
            {
                managedObject = null;
                return false;
            }
 
            managedObject = _managedObjectWeakReference.Target as TObject;
            return managedObject != null;
        }
 
        public ComHandle<THandle, TObject>? ComHandle
        {
            get
            {
                var rcw = this.ComAggregateObject;
                if (rcw == null)
                {
                    return null;
                }
 
                Debug.Assert(_managedObjectWeakReference != null);
                if (_managedObjectWeakReference.Target is TObject managedObject)
                {
                    // Construct a new ComHandle without going through the cycle of unwrapping
                    // the managed object from the rcw, that has shown to be a perf concern for 
                    // progression (see Dev10 Bug 628992).
                    return new ComHandle<THandle, TObject>(rcw, managedObject);
                }
                else
                {
                    // We fall back to trying to unwrap the managed object out of the rcw, but
                    // the Weakref to the managed object shouldn't go null if the rcw is still
                    // alive, should it?
                    Debug.Fail("Can this really happen?");
                    return new ComHandle<THandle, TObject>(rcw);
                }
            }
        }
 
        public bool IsAlive()
        {
            if (_managedObjectWeakReference == null)
            {
                return false;
            }
 
            return _managedObjectWeakReference.IsAlive;
        }
    }
}