// 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.Linq;
using System.Threading;
using System.Threading.Tasks;
 
namespace Microsoft.CodeAnalysis.Editing
{
    public static class SymbolEditorExtensions
    {
        /// <summary>
        /// Gets the reference to the declaration of the base or interface type as part of the symbol's declaration. 
        /// </summary>
        public static async Task<SyntaxNode> GetBaseOrInterfaceDeclarationReferenceAsync(
            this SymbolEditor editor,
            ISymbol symbol,
            ITypeSymbol baseOrInterfaceType,
            CancellationToken cancellationToken = default)
        {
            if (baseOrInterfaceType == null)
            {
                throw new ArgumentNullException(nameof(baseOrInterfaceType));
            }
 
            if (baseOrInterfaceType.TypeKind != TypeKind.Error)
            {
                baseOrInterfaceType = (ITypeSymbol)(await editor.GetCurrentSymbolAsync(baseOrInterfaceType, cancellationToken).ConfigureAwait(false));
            }
 
            // look for the base or interface declaration in all declarations of the symbol
            var currentDecls = await editor.GetCurrentDeclarationsAsync(symbol, cancellationToken).ConfigureAwait(false);
 
            foreach (var decl in currentDecls)
            {
                var doc = editor.OriginalSolution.GetDocument(decl.SyntaxTree);
                var model = await doc.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
                var gen = SyntaxGenerator.GetGenerator(doc);
 
                var typeRef = gen.GetBaseAndInterfaceTypes(decl).FirstOrDefault(r => model.GetTypeInfo(r, cancellationToken).Type.Equals(baseOrInterfaceType));
                if (typeRef != null)
                {
                    return typeRef;
                }
            }
 
            return null;
        }
 
        /// <summary>
        /// Changes the base type of the symbol.
        /// </summary>
        public static async Task<ISymbol> SetBaseTypeAsync(
            this SymbolEditor editor,
            INamedTypeSymbol symbol,
            Func<SyntaxGenerator, SyntaxNode> getNewBaseType,
            CancellationToken cancellationToken = default)
        {
            var baseType = symbol.BaseType;
 
            if (baseType != null)
            {
                // find existing declaration of the base type
                var typeRef = await editor.GetBaseOrInterfaceDeclarationReferenceAsync(symbol, baseType, cancellationToken).ConfigureAwait(false);
                if (typeRef != null)
                {
                    return await editor.EditOneDeclarationAsync(
                        symbol,
                        typeRef.GetLocation(),
                        (e, d) => e.ReplaceNode(typeRef, getNewBaseType(e.Generator)),
                        cancellationToken).ConfigureAwait(false);
                }
            }
 
            // couldn't find the existing reference to change, so add it to one of the declarations
            return await editor.EditOneDeclarationAsync(symbol, (e, decl) =>
            {
                var newBaseType = getNewBaseType(e.Generator);
                if (newBaseType != null)
                {
                    e.ReplaceNode(decl, (d, g) => g.AddBaseType(d, newBaseType));
                }
            }, cancellationToken).ConfigureAwait(false);
        }
 
        /// <summary>
        /// Changes the base type of the symbol.
        /// </summary>
        public static Task<ISymbol> SetBaseTypeAsync(
            this SymbolEditor editor,
            INamedTypeSymbol symbol,
            ITypeSymbol newBaseType,
            CancellationToken cancellationToken = default)
        {
            return editor.SetBaseTypeAsync(symbol, g => newBaseType != null ? g.TypeExpression(newBaseType) : null, cancellationToken);
        }
    }
}