// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System.Collections.Generic;
using System.Diagnostics;
using ILLink.Shared.TypeSystemProxy;
using Mono.Cecil;

namespace Mono.Linker.Steps
{
    public class DiscoverOperatorsHandler : IMarkHandler
    {
        LinkContext? _context;
        LinkContext Context
        {
            get
            {
                Debug.Assert(_context != null);
                return _context;
            }
        }

        bool _seenLinqExpressions;
        readonly HashSet<TypeDefinition> _trackedTypesWithOperators;
        Dictionary<TypeDefinition, List<MethodDefinition>>? _pendingOperatorsForType;

        Dictionary<TypeDefinition, List<MethodDefinition>> PendingOperatorsForType
        {
            get
            {
                _pendingOperatorsForType ??= new Dictionary<TypeDefinition, List<MethodDefinition>>();
                return _pendingOperatorsForType;
            }
        }

        public DiscoverOperatorsHandler()
        {
            _trackedTypesWithOperators = new HashSet<TypeDefinition>();
        }

        public void Initialize(LinkContext context, MarkContext markContext)
        {
            _context = context;
            markContext.RegisterMarkTypeAction(ProcessType);
        }

        void ProcessType(TypeDefinition type)
        {
            CheckForLinqExpressions(type);

            // Check for custom operators and either:
            // - mark them, if Linq.Expressions was already marked, or
            // - track them to be marked in case Linq.Expressions is marked later
            var hasOperators = ProcessCustomOperators(type, mark: _seenLinqExpressions);
            if (!_seenLinqExpressions)
            {
                if (hasOperators)
                    _trackedTypesWithOperators.Add(type);
                return;
            }

            // Mark pending operators defined on other types that reference this type
            // (these are only tracked if we have already seen Linq.Expressions)
            if (PendingOperatorsForType.TryGetValue(type, out var pendingOperators))
            {
                foreach (var customOperator in pendingOperators)
                    MarkOperator(customOperator);
                PendingOperatorsForType.Remove(type);
            }
        }

        void CheckForLinqExpressions(TypeDefinition type)
        {
            if (_seenLinqExpressions)
                return;

            if (type.Namespace != "System.Linq.Expressions" || type.Name != "Expression")
                return;

            _seenLinqExpressions = true;

            foreach (var markedType in _trackedTypesWithOperators)
                ProcessCustomOperators(markedType, mark: true);

            _trackedTypesWithOperators.Clear();
        }

        void MarkOperator(MethodDefinition method)
        {
            Context.Annotations.Mark(method, new DependencyInfo(DependencyKind.PreservedOperator, method.DeclaringType), new MessageOrigin(method.DeclaringType));
        }

        bool ProcessCustomOperators(TypeDefinition type, bool mark)
        {
            if (!type.HasMethods)
                return false;

            bool hasCustomOperators = false;
            foreach (var method in type.Methods)
            {
                if (!IsOperator(method, out var otherType))
                    continue;

                if (!mark)
                    return true;

                Debug.Assert(_seenLinqExpressions);
                hasCustomOperators = true;

                if (otherType == null || Context.Annotations.IsMarked(otherType))
                {
                    MarkOperator(method);
                    continue;
                }

                // Wait until otherType gets marked to mark the operator.
                if (!PendingOperatorsForType.TryGetValue(otherType, out var pendingOperators))
                {
                    pendingOperators = new List<MethodDefinition>();
                    PendingOperatorsForType.Add(otherType, pendingOperators);
                }
                pendingOperators.Add(method);
            }
            return hasCustomOperators;
        }

        TypeDefinition? _nullableOfT;
        TypeDefinition? NullableOfT
        {
            get
            {
                _nullableOfT ??= BCL.FindPredefinedType(WellKnownType.System_Nullable_T, Context);
                return _nullableOfT;
            }
        }

        TypeDefinition? NonNullableType(TypeReference type)
        {
            var typeDef = Context.TryResolve(type);
            if (typeDef == null)
                return null;

            if (!typeDef.IsValueType || typeDef != NullableOfT)
                return typeDef;

            // Unwrap Nullable<T>
            Debug.Assert(typeDef.HasGenericParameters);
            // The original type reference might be a TypeSpecification like array of Nullable<T>
            // that we need to unwrap until we get to the Nullable<T>
            while (!type.IsGenericInstance)
                type = ((TypeSpecification)type).ElementType;
            var nullableType = type as GenericInstanceType;
            Debug.Assert(nullableType != null && nullableType.HasGenericArguments && nullableType.GenericArguments.Count == 1);
            return Context.TryResolve(nullableType.GenericArguments[0]);
        }

        bool IsOperator(MethodDefinition method, out TypeDefinition? otherType)
        {
            otherType = null;

            if (!method.IsStatic || !method.IsPublic || !method.IsSpecialName || !method.Name.StartsWith("op_"))
                return false;

            var operatorName = method.Name.Substring(3);
            var self = method.DeclaringType;

            switch (operatorName)
            {
                // Unary operators
                case "UnaryPlus":
                case "UnaryNegation":
                case "LogicalNot":
                case "OnesComplement":
                case "Increment":
                case "Decrement":
                case "True":
                case "False":
                    // Parameter type of a unary operator must be the declaring type
                    if (method.GetMetadataParametersCount() != 1 || NonNullableType(method.GetParameter((ParameterIndex)0).ParameterType) != self)
                        return false;
                    // ++ and -- must return the declaring type
                    if (operatorName is "Increment" or "Decrement" && NonNullableType(method.ReturnType) != self)
                        return false;
                    return true;
                // Binary operators
                case "Addition":
                case "Subtraction":
                case "Multiply":
                case "Division":
                case "Modulus":
                case "BitwiseAnd":
                case "BitwiseOr":
                case "ExclusiveOr":
                case "LeftShift":
                case "RightShift":
                case "Equality":
                case "Inequality":
                case "LessThan":
                case "GreaterThan":
                case "LessThanOrEqual":
                case "GreaterThanOrEqual":
                    if (method.GetMetadataParametersCount() != 2)
                        return false;
                    var nnLeft = NonNullableType(method.GetParameter((ParameterIndex)0).ParameterType);
                    var nnRight = NonNullableType(method.GetParameter((ParameterIndex)1).ParameterType);
                    if (nnLeft == null || nnRight == null)
                        return false;
                    // << and >> must take the declaring type and int
                    if (operatorName is "LeftShift" or "RightShift" && (nnLeft != self || nnRight.MetadataType != MetadataType.Int32))
                        return false;
                    // At least one argument must be the declaring type
                    if (nnLeft != self && nnRight != self)
                        return false;
                    if (nnLeft != self)
                        otherType = nnLeft;
                    if (nnRight != self)
                        otherType = nnRight;
                    return true;
                // Conversion operators
                case "Implicit":
                case "Explicit":
                    if (method.GetMetadataParametersCount() != 1)
                        return false;
                    var nnSource = NonNullableType(method.GetParameter((ParameterIndex)0).ParameterType);
                    var nnTarget = NonNullableType(method.ReturnType);
                    // Exactly one of source/target must be the declaring type
                    if (nnSource == self == (nnTarget == self))
                        return false;
                    otherType = nnSource == self ? nnTarget : nnSource;
                    return true;
                default:
                    return false;
            }
        }
    }
}