using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.FindSymbols;
using Microsoft.CodeAnalysis.Operations;
using Microsoft.CodeAnalysis.Text;
using Microsoft.Extensions.Options;

namespace Cfa;

internal sealed class MessageLinks
{
    private readonly MessageOptions options;
    private readonly IReadOnlyCollection<Compilation> compilations;
    private readonly HashSet<SyntaxTree> includedTrees;
    private readonly Solution solution;
    private readonly Dictionary<(MessageRoute Route, string Type, bool Publishers), HashSet<ISymbol>> links = [];
    private readonly Dictionary<(SyntaxTree Tree, TextSpan Span), Registration[]> registrations = [];
    private readonly Dictionary<(MethodSelector Method, int Argument, int? TypeArgument), IMethodSymbol[]> methods = [];
    private readonly IImmutableSet<Document> documents;

    public MessageLinks(
        MessageOptions options,
        IReadOnlyCollection<Compilation> compilations,
        HashSet<SyntaxTree> includedTrees,
        Solution solution
    )
    {
        this.options = options;
        this.compilations = compilations;
        this.includedTrees = includedTrees;
        this.solution = solution;
        documents = includedTrees.Select(solution.GetDocument).OfType<Document>().ToImmutableHashSet();
        // Validate API shapes only; call sites are resolved on demand.
        foreach (var route in options.MessageRoutes)
        {
            Validate(route.Publisher.Method, route.Publisher.MessageArgument);
            foreach (var rule in route.Registrations)
            {
                Validate(rule.Method, rule.HandlerArgument, rule.MessageTypeArgument);
            }
        }
    }

    private sealed record Registration(MessageRoute Route, ITypeSymbol Message, ISymbol Handler, TextSpan Argument);

    public bool IsHandler(ISymbol member) =>
        options.MessageRoutes is not []
        && member is IMethodSymbol { MethodKind: MethodKind.AnonymousFunction }
        && member.DeclaringSyntaxReferences.Any(reference =>
            reference
                .GetSyntax()
                .Ancestors()
                .OfType<InvocationExpressionSyntax>()
                .Any(invocation =>
                    Registrations(invocation)
                        .Any(registration => SymbolEqualityComparer.Default.Equals(registration.Handler, member))
                )
        );

    public bool IsRegistrationArgument(SyntaxNode node) =>
        options.MessageRoutes is not []
        && node.AncestorsAndSelf()
            .TakeWhile(ancestor => ancestor is not AnonymousFunctionExpressionSyntax)
            .OfType<ArgumentSyntax>()
            .Any(argument =>
                argument.Parent?.Parent is InvocationExpressionSyntax invocation
                && Registrations(invocation).Any(registration => registration.Argument == argument.Span)
            );

    public async Task<IEnumerable<ISymbol>> Targets(ISymbol member, CancellationToken token)
    {
        var result = new HashSet<ISymbol>(SymbolEqualityComparer.Default);
        if (options.MessageRoutes is [])
        {
            return result;
        }
        foreach (var reference in member.DeclaringSyntaxReferences)
        {
            var syntax = await reference.GetSyntaxAsync(token);
            var model = compilations.GetSemanticModel(syntax.SyntaxTree);
            foreach (
                var invocation in MemberSyntax
                    .ExecutionRoots(member, syntax)
                    .SelectMany(root =>
                        root.DescendantNodesAndSelf(node =>
                            node == root
                            || node is not AnonymousFunctionExpressionSyntax
                            || model.GetSymbolInfo(node, token).Symbol is not { } handler
                            || !IsHandler(handler)
                        )
                    )
                    .OfType<InvocationExpressionSyntax>()
                    .Where(invocation =>
                        options.MessageRoutes.Any(route => route.Publisher.Method.Name == InvocationName(invocation))
                    )
            )
            {
                token.ThrowIfCancellationRequested();
                if (
                    !(member is IMethodSymbol { MethodKind: MethodKind.Constructor } && syntax is TypeDeclarationSyntax)
                        && !SameMember(ContainingMember(model, invocation, token), member)
                    || model.GetOperation(invocation, token) is not IInvocationOperation operation
                )
                {
                    continue;
                }
                foreach (
                    var route in options.MessageRoutes.Where(route =>
                        Matches(operation.TargetMethod, route.Publisher.Method)
                    )
                )
                {
                    if (PublishedMessage(operation, route) is { } message)
                    {
                        result.UnionWith(await Resolve(route, message, publishers: false, token));
                    }
                }
            }
        }
        return result;
    }

    public async Task<IEnumerable<ISymbol>> Callers(
        ISymbol member,
        IEnumerable<ReferencedSymbol> references,
        CancellationToken token
    )
    {
        var result = new HashSet<ISymbol>(SymbolEqualityComparer.Default);
        if (options.MessageRoutes is [])
        {
            return result;
        }
        var candidates = new HashSet<InvocationExpressionSyntax>();
        if (member is IMethodSymbol { MethodKind: MethodKind.AnonymousFunction })
        {
            foreach (var reference in member.DeclaringSyntaxReferences)
            {
                candidates.UnionWith(
                    (await reference.GetSyntaxAsync(token)).Ancestors().OfType<InvocationExpressionSyntax>()
                );
            }
        }
        else
        {
            candidates.UnionWith(await ReferencingInvocations(references, token));
        }
        foreach (
            var registration in candidates
                .SelectMany(Registrations)
                .Where(registration => SameMember(registration.Handler, member))
        )
        {
            result.UnionWith(await Resolve(registration.Route, registration.Message, publishers: true, token));
        }
        return result;
    }

    private Registration[] Registrations(InvocationExpressionSyntax syntax)
    {
        var key = (syntax.SyntaxTree, syntax.Span);
        if (registrations.TryGetValue(key, out var found))
        {
            return found;
        }
        var result = new List<Registration>();
        if (!includedTrees.Contains(syntax.SyntaxTree))
        {
            return [];
        }
        var name = InvocationName(syntax);
        foreach (var route in options.MessageRoutes)
        {
            foreach (var rule in route.Registrations.Where(rule => rule.Method.Name == name))
            {
                var model = compilations.GetSemanticModel(syntax.SyntaxTree);
                if (
                    model.GetOperation(syntax) is not IInvocationOperation operation
                    || !Matches(operation.TargetMethod, rule.Method)
                )
                {
                    continue;
                }
                Validate(rule.Method, rule.HandlerArgument, rule.MessageTypeArgument);
                var callback = operation.Arguments.FirstOrDefault(argument =>
                    argument.Parameter?.Ordinal == rule.HandlerArgument
                );
                if (callback is not { Value: var callbackOperation })
                {
                    continue;
                }
                var handler = Unwrap(callbackOperation) switch
                {
                    IAnonymousFunctionOperation lambda => lambda.Symbol,
                    IMethodReferenceOperation method => method.Method,
                    _ => null,
                };
                if (
                    handler is null
                    || !handler.DeclaringSyntaxReferences.Any(reference => includedTrees.Contains(reference.SyntaxTree))
                )
                {
                    continue;
                }
                var argument = syntax.ArgumentList.Arguments.First(argument =>
                    argument.Span.Contains(callbackOperation.Syntax.Span)
                );
                result.Add(
                    new(route, operation.TargetMethod.TypeArguments[rule.MessageTypeArgument], handler, argument.Span)
                );
            }
        }
        found = [.. result];
        registrations.Add(key, found);
        return found;
    }

    private static string InvocationName(InvocationExpressionSyntax invocation) =>
        invocation.Expression switch
        {
            SimpleNameSyntax simple => simple.Identifier.ValueText,
            MemberAccessExpressionSyntax access => access.Name.Identifier.ValueText,
            MemberBindingExpressionSyntax binding => binding.Name.Identifier.ValueText,
            _ => "",
        };

    private ITypeSymbol? PublishedMessage(IInvocationOperation operation, MessageRoute route)
    {
        Validate(route.Publisher.Method, route.Publisher.MessageArgument);
        var message = Unwrap(
            operation
                .Arguments.FirstOrDefault(argument => argument.Parameter?.Ordinal == route.Publisher.MessageArgument)
                ?.Value
        )?.Type;
        return message is INamedTypeSymbol { IsAbstract: false, TypeKind: TypeKind.Class or TypeKind.Struct }
            ? message
            : null;
    }

    private async Task<HashSet<ISymbol>> Resolve(
        MessageRoute route,
        ITypeSymbol message,
        bool publishers,
        CancellationToken token
    )
    {
        var key = (
            route,
            $"{message.ContainingAssembly.Identity}:{DocumentationCommentId.CreateReferenceId(message)}",
            publishers
        );
        if (links.TryGetValue(key, out var found))
        {
            return found;
        }
        var result = new HashSet<ISymbol>(SymbolEqualityComparer.Default);
        if (publishers)
        {
            foreach (var method in Validate(route.Publisher.Method, route.Publisher.MessageArgument))
            {
                foreach (var invocation in await ReferencingInvocations(method, token))
                {
                    var model = compilations.GetSemanticModel(invocation.SyntaxTree);
                    if (
                        model.GetOperation(invocation, token) is not IInvocationOperation operation
                        || !Matches(operation.TargetMethod, route.Publisher.Method)
                        || PublishedMessage(operation, route) is not { } published
                        || !SameMessageType(published, message)
                    )
                    {
                        continue;
                    }
                    var caller = ContainingMember(model, invocation, token);
                    if (caller is IMethodSymbol or IPropertySymbol)
                    {
                        result.Add(caller);
                    }
                }
            }
        }
        else
        {
            foreach (var rule in route.Registrations)
            {
                Validate(rule.Method, rule.HandlerArgument, rule.MessageTypeArgument);
            }
            // Search references to this message type, rather than every registration API call.
            foreach (var invocation in await ReferencingInvocations(message.OriginalDefinition, token))
            {
                foreach (var registration in Registrations(invocation))
                {
                    if (registration.Route == route && SameMessageType(message, registration.Message))
                    {
                        result.Add(registration.Handler);
                    }
                }
            }
        }
        links.Add(key, result);
        return result;
    }

    private ISymbol? ContainingMember(
        SemanticModel model,
        InvocationExpressionSyntax invocation,
        CancellationToken token
    )
    {
        var caller = model.GetEnclosingSymbol(invocation.SpanStart, token);
        while (caller is IMethodSymbol { MethodKind: MethodKind.AnonymousFunction } && !IsHandler(caller))
        {
            caller = caller.ContainingSymbol;
        }
        return caller is IMethodSymbol { AssociatedSymbol: IPropertySymbol property } ? property : caller;
    }

    private async Task<HashSet<InvocationExpressionSyntax>> ReferencingInvocations(
        ISymbol symbol,
        CancellationToken token
    ) =>
        await ReferencingInvocations(await SymbolFinder.FindReferencesAsync(symbol, solution, documents, token), token);

    private static async Task<HashSet<InvocationExpressionSyntax>> ReferencingInvocations(
        IEnumerable<ReferencedSymbol> references,
        CancellationToken token
    )
    {
        var result = new HashSet<InvocationExpressionSyntax>();
        foreach (var reference in references)
        {
            foreach (var location in reference.Locations)
            {
                token.ThrowIfCancellationRequested();
                var root = await location.Document.GetSyntaxRootAsync(token);
                result.UnionWith(
                    root!.FindNode(location.Location.SourceSpan).AncestorsAndSelf().OfType<InvocationExpressionSyntax>()
                );
            }
        }
        return result;
    }

    private static bool SameMessageType(ITypeSymbol left, ITypeSymbol right) =>
        SymbolEqualityComparer.Default.Equals(left, right)
        // A referenced project's type can be represented by a metadata symbol.
        || left.ContainingAssembly.Identity.Equals(right.ContainingAssembly.Identity)
            && DocumentationCommentId.CreateReferenceId(left) is { } id
            && id == DocumentationCommentId.CreateReferenceId(right);

    private static bool SameMember(ISymbol? left, ISymbol right) =>
        SymbolEqualityComparer.Default.Equals(left?.OriginalDefinition, right.OriginalDefinition)
        // Retargeted symbols from another project still represent the same source declaration.
        || left is not null
            && left.ContainingAssembly.Identity.Equals(right.ContainingAssembly.Identity)
            && left.OriginalDefinition.GetDocumentationCommentId() is { } id
            && id == right.OriginalDefinition.GetDocumentationCommentId();

    private static IOperation? Unwrap(IOperation? operation) =>
        operation switch
        {
            IConversionOperation { IsImplicit: true } conversion => Unwrap(conversion.Operand),
            IDelegateCreationOperation creation => Unwrap(creation.Target),
            _ => operation,
        };

    private static bool Matches(IMethodSymbol method, MethodSelector selector)
    {
        for (var member = method; member is not null; member = member.OverriddenMethod)
        {
            if (MatchesDefinition(member, selector))
            {
                return true;
            }
        }
        foreach (
            var contract in method.ContainingType.AllInterfaces.Where(contract =>
                MetadataName(contract) == selector.Type
            )
        )
        {
            foreach (
                var member in contract
                    .GetMembers(selector.Name)
                    .OfType<IMethodSymbol>()
                    .Where(member => MatchesDefinition(member, selector))
            )
            {
                if (
                    SymbolEqualityComparer.Default.Equals(
                        method.ContainingType.FindImplementationForInterfaceMember(member)?.OriginalDefinition,
                        method.OriginalDefinition
                    )
                )
                {
                    return true;
                }
            }
        }
        return false;
    }

    private static bool MatchesDefinition(IMethodSymbol method, MethodSelector selector) =>
        method.Name == selector.Name
        && MetadataName(method.ContainingType) == selector.Type
        && (selector.GenericArity is null || method.Arity == selector.GenericArity)
        && (selector.ParameterCount is null || method.Parameters.Length == selector.ParameterCount);

    private static string MetadataName(INamedTypeSymbol type) =>
        type.ContainingType is { } parent ? $"{MetadataName(parent)}+{type.MetadataName}"
        : type.ContainingNamespace.IsGlobalNamespace ? type.MetadataName
        : $"{type.ContainingNamespace}.{type.MetadataName}";

    private IMethodSymbol[] Validate(MethodSelector selector, int argument, int? typeArgument = null)
    {
        var key = (selector, argument, typeArgument);
        if (methods.TryGetValue(key, out var found))
        {
            return found;
        }
        found =
        [
            .. compilations
                .Select(compilation => compilation.GetTypeByMetadataName(selector.Type))
                .OfType<INamedTypeSymbol>()
                .SelectMany(type => type.GetMembers(selector.Name))
                .OfType<IMethodSymbol>()
                .Where(method => MatchesDefinition(method, selector))
                .GroupBy(method => (method.ContainingAssembly.Identity, method.GetDocumentationCommentId()))
                .Select(group =>
                    group.FirstOrDefault(method => method.DeclaringSyntaxReferences is not []) ?? group.First()
                ),
        ];
        if (
            found is []
            || found.Any(method =>
                argument < 0
                || argument >= method.Parameters.Length
                || typeArgument is { } index && (index < 0 || index >= method.Arity)
                || typeArgument is not null && method.Parameters[argument].Type.TypeKind != TypeKind.Delegate
            )
        )
        {
            throw new OptionsValidationException(
                nameof(MessageOptions),
                typeof(MessageOptions),
                [
                    $"Message method '{selector.Type}.{selector.Name}' was not found or its configured argument indices are invalid.",
                ]
            );
        }
        methods.Add(key, found);
        return found;
    }
}