// 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;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using Mono.Cecil;
 
namespace Mono.Linker
{
	public class CustomAttributeSource
	{
		public AttributeInfo PrimaryAttributeInfo { get; }
		private readonly Dictionary<AssemblyDefinition, AttributeInfo?> _embeddedXmlInfos;
		readonly LinkContext _context;
 
		public CustomAttributeSource (LinkContext context)
		{
			PrimaryAttributeInfo = new AttributeInfo ();
			_embeddedXmlInfos = new Dictionary<AssemblyDefinition, AttributeInfo?> ();
			_context = context;
		}
 
		public static AssemblyDefinition GetAssemblyFromCustomAttributeProvider (ICustomAttributeProvider provider)
		{
			return provider switch {
				MemberReference mr => mr.Module.Assembly,
				AssemblyDefinition ad => ad,
				ModuleDefinition md => md.Assembly,
				InterfaceImplementation ii => ii.InterfaceType.Module.Assembly,
				GenericParameterConstraint gpc => gpc.ConstraintType.Module.Assembly,
				ParameterDefinition pd => pd.ParameterType.Module.Assembly,
				MethodReturnType mrt => mrt.ReturnType.Module.Assembly,
				_ => throw new NotImplementedException (provider.GetType ().ToString ()),
			};
		}
 
		public bool TryGetEmbeddedXmlInfo (ICustomAttributeProvider provider, [NotNullWhen (true)] out AttributeInfo? xmlInfo)
		{
			AssemblyDefinition assembly = GetAssemblyFromCustomAttributeProvider (provider);
 
			if (!_embeddedXmlInfos.TryGetValue (assembly, out xmlInfo)) {
				// Add an empty record - this prevents reentrancy
				// If the embedded XML itself generates warnings, trying to log such warning
				// may ask for attributes (suppressions) and thus we could end up in this very place again
				// So first add a dummy record and once processed we will replace it with the real data
				_embeddedXmlInfos.Add (assembly, new AttributeInfo ());
				xmlInfo = _context.EmbeddedXmlInfo.ProcessAttributes (assembly, _context);
				_embeddedXmlInfos[assembly] = xmlInfo;
			}
 
			return xmlInfo != null;
		}
 
		public IEnumerable<CustomAttribute> GetCustomAttributes (ICustomAttributeProvider provider, string attributeNamespace, string attributeName)
		{
			foreach (var attr in GetCustomAttributes (provider)) {
				if (attr.AttributeType.Namespace == attributeNamespace && attr.AttributeType.Name == attributeName)
					yield return attr;
			}
		}
 
		public IEnumerable<CustomAttribute> GetCustomAttributes (ICustomAttributeProvider provider)
		{
			if (provider.HasCustomAttributes) {
				foreach (var customAttribute in provider.CustomAttributes)
					yield return customAttribute;
			}
 
			if (PrimaryAttributeInfo.CustomAttributes.TryGetValue (provider, out var annotations)) {
				foreach (var customAttribute in annotations)
					yield return customAttribute;
			}
 
			if (!TryGetEmbeddedXmlInfo (provider, out var embeddedXml))
				yield break;
 
			if (embeddedXml.CustomAttributes.TryGetValue (provider, out annotations)) {
				foreach (var customAttribute in annotations)
					yield return customAttribute;
			}
		}
 
		public bool TryGetCustomAttributeOrigin (ICustomAttributeProvider provider, CustomAttribute customAttribute, out MessageOrigin origin)
		{
			if (PrimaryAttributeInfo.CustomAttributesOrigins.TryGetValue (customAttribute, out origin))
				return true;
 
			if (!TryGetEmbeddedXmlInfo (provider, out var embeddedXml))
				return false;
 
			return embeddedXml.CustomAttributesOrigins.TryGetValue (customAttribute, out origin);
		}
 
		public bool HasAny (ICustomAttributeProvider provider)
		{
			if (provider.HasCustomAttributes)
				return true;
 
			if (PrimaryAttributeInfo.CustomAttributes.ContainsKey (provider))
				return true;
 
			if (!TryGetEmbeddedXmlInfo (provider, out var embeddedXml))
				return false;
 
			return embeddedXml.CustomAttributes.ContainsKey (provider);
		}
	}
}