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

using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;

#if FEATURE_COMINTEROP || FEATURE_COMWRAPPERS

using static System.WeakReferenceHandleTags;

namespace System
{
    internal sealed partial class ComAwareWeakReference
    {
        // _weakHandle is effectively readonly.
        // the only place where we change it after construction is in finalizer.
        private nint _weakHandle;
        private ComInfo? _comInfo;

        internal sealed class ComInfo
        {
            // _pComWeakRef is effectively readonly.
            // the only place where we change it after construction is in finalizer.
            private IntPtr _pComWeakRef;
            private readonly object? _context;

            internal object? ResolveTarget()
            {
                return ComWeakRefToObject(_pComWeakRef, _context);
            }

            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            internal static ComInfo? FromObject(object? target)
            {
                if (target == null || !PossiblyComObject(target))
                    return null;

                return FromObjectSlow(target);
            }

            [MethodImpl(MethodImplOptions.NoInlining)]
            private static ComInfo? FromObjectSlow(object target)
            {
                IntPtr pComWeakRef = ObjectToComWeakRef(target, out object? context);
                if (pComWeakRef == 0)
                    return null;

                try
                {
                    return new ComInfo(pComWeakRef, context);
                }
                catch (OutOfMemoryException)
                {
                    // we did not create an object, so ComInfo finalizer will not run
                    Marshal.Release(pComWeakRef);
                    throw;
                }
            }

            private ComInfo(IntPtr pComWeakRef, object? context)
            {
                Debug.Assert(pComWeakRef != IntPtr.Zero);
                _pComWeakRef = pComWeakRef;
                _context = context;
            }

            ~ComInfo()
            {
                Marshal.Release(_pComWeakRef);
                // our use pattern guarantees that the instance is not reachable after this.
                // clear the pointer just in case that gets somehow broken.
                _pComWeakRef = 0;
            }
        }

        private ComAwareWeakReference(nint weakHandle)
        {
            Debug.Assert(weakHandle != 0);
            _weakHandle = weakHandle;
        }

        ~ComAwareWeakReference()
        {
            GCHandle.InternalFree(_weakHandle);
            // our use pattern guarantees that the instance is not reachable after this.
            // clear the pointer just in case that gets somehow broken.
            _weakHandle = 0;
        }

        private void SetTarget(object? target, ComInfo? comInfo)
        {
            // NOTE: ComAwareWeakReference is an internal implementation detail and
            //       instances are never exposed publicly, thus we can use "this" for locking
            lock (this)
            {
                GCHandle.InternalSet(_weakHandle, target);
                _comInfo = comInfo;
            }
        }

        internal object? Target => GCHandle.InternalGet(_weakHandle);

        internal nint WeakHandle => _weakHandle;

        internal T? RehydrateTarget<T>() where T: class?
        {
            T? target = null;
            lock (this)
            {
                if (_comInfo != null)
                {
                    // Check if the target is still null
                    target = Unsafe.As<T>(GCHandle.InternalGet(_weakHandle));
                    if (target == null)
                    {
                        // Resolve and reset. Perform runtime cast to catch bugs
                        // in COM interop where it rehydrates wrong type.
                        target = (T?)_comInfo.ResolveTarget();
                        if (target != null)
                            GCHandle.InternalSet(_weakHandle, target);
                    }
                }
            }

            return target;
        }

        private static ComAwareWeakReference EnsureComAwareReference(ref nint taggedHandle)
        {
            nint current = taggedHandle;
            if ((current & ComAwareBit) == 0)
            {
                ComAwareWeakReference newRef = new ComAwareWeakReference(taggedHandle & ~HandleTagBits);
                nint newHandle = GCHandle.InternalAlloc(newRef, GCHandleType.Normal);
                nint newTaggedHandle = newHandle | ComAwareBit | (taggedHandle & TracksResurrectionBit);
                if (Interlocked.CompareExchange(ref taggedHandle, newTaggedHandle, current) == current)
                {
                    // success.
                    return newRef;
                }

                // someone beat us to it. (this is rare)
                GCHandle.InternalFree(newHandle);
                GC.SuppressFinalize(newRef);
            }

            return Unsafe.As<ComAwareWeakReference>(GCHandle.InternalGet(taggedHandle & ~HandleTagBits)!);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal static ComAwareWeakReference GetFromTaggedReference(nint taggedHandle)
        {
            Debug.Assert((taggedHandle & ComAwareBit) != 0);
            return Unsafe.As<ComAwareWeakReference>(GCHandle.InternalGet(taggedHandle & ~HandleTagBits)!);
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        internal static void SetTarget(ref nint taggedHandle, object? target, ComInfo? comInfo)
        {
            ComAwareWeakReference comAwareRef = comInfo != null ?
                EnsureComAwareReference(ref taggedHandle) :
                Unsafe.As<ComAwareWeakReference>(GCHandle.InternalGet(taggedHandle & ~HandleTagBits)!);

            comAwareRef.SetTarget(target, comInfo);
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        internal static void SetComInfoInConstructor(ref nint taggedHandle, ComInfo comInfo)
        {
            Debug.Assert((taggedHandle & ComAwareBit) == 0);
            ComAwareWeakReference comAwareRef = new ComAwareWeakReference(taggedHandle & ~HandleTagBits);
            nint newHandle = GCHandle.InternalAlloc(comAwareRef, GCHandleType.Normal);
            taggedHandle = newHandle | ComAwareBit | (taggedHandle & TracksResurrectionBit);
            comAwareRef._comInfo = comInfo;
        }
    }
}
#endif