using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace Cfa;

internal sealed class CalleeBuilder
{
    public static IImmutableDictionary<ISymbol, IImmutableList<Call>> FindCallees(
        IReadOnlyCollection<Compilation> compilations,
        ISymbol symbol,
        CancellationToken cancellationToken
    )
    {
        var builder = new CalleeBuilder(
            compilations,
            [
                .. compilations
                    .SelectMany(compilation =>
                        compilation.GetSymbolsWithName(_ => true, SymbolFilter.Type, cancellationToken)
                    )
                    .OfType<INamedTypeSymbol>(),
            ]
        );

        builder.FindCallees(symbol, cancellationToken);
        return builder.reachableFromSymbols.ToImmutable();
    }

    private readonly IReadOnlyCollection<Compilation> compilations;
    private readonly IReadOnlyCollection<INamedTypeSymbol> types;
    private readonly HashSet<ISymbol> seenMember = [with(SymbolEqualityComparer.Default)];
    private readonly ImmutableDictionary<ISymbol, IImmutableList<Call>>.Builder reachableFromSymbols =
        ImmutableDictionary.CreateBuilder<ISymbol, IImmutableList<Call>>(SymbolEqualityComparer.Default);

    private CalleeBuilder(IReadOnlyCollection<Compilation> compilations, IReadOnlyCollection<INamedTypeSymbol> types)
    {
        this.compilations = compilations;
        this.types = types;
    }

    private void FindCallees(ISymbol symbol, CancellationToken cancellationToken)
    {
        if (!seenMember.Add(symbol))
        {
            return;
        }

        var seenCall = new Dictionary<ISymbol, LinkKind>(SymbolEqualityComparer.Default);
        foreach (var reference in symbol.DeclaringSyntaxReferences)
        {
            var syntax = reference.GetSyntax(cancellationToken);
            var semanticModelForSyntax = compilations.GetSemanticModel(syntax.SyntaxTree);
            var accesses = syntax
                .DescendantNodesAndSelf()
                .Where(node =>
                    node
                        is SimpleNameSyntax
                            or ElementAccessExpressionSyntax
                            or ElementBindingExpressionSyntax
                            or BaseObjectCreationExpressionSyntax
                            or ConstructorInitializerSyntax
                );
            foreach (var access in accesses)
            {
                var accessedSymbol = semanticModelForSyntax.GetSymbolInfo(access, cancellationToken).Symbol;
                if (accessedSymbol is not (IMethodSymbol or IPropertySymbol))
                {
                    continue;
                }

                if (accessedSymbol.IsAbstract || accessedSymbol.IsVirtual)
                {
                    foreach (var implementation in FindImplementations(accessedSymbol))
                    {
                        AddCall(implementation, LinkKind.Abstract);
                    }
                }

                if (!accessedSymbol.IsAbstract)
                {
                    AddCall(accessedSymbol, LinkKind.Direct);
                }
            }
        }
        reachableFromSymbols.Add(symbol, seenCall.Select(call => new Call(call.Key, call.Value)).ToImmutableList());

        void AddCall(ISymbol target, LinkKind linkKind)
        {
            if (
                !IsDeclaredInSource(target)
                || linkKind == LinkKind.Abstract && target.IsAbstract
                || target is IPropertySymbol property && !HasBody(property, cancellationToken)
            )
            {
                return;
            }

            if (seenCall.TryAdd(target, linkKind))
            {
                FindCallees(target, cancellationToken);
            }
            else if (linkKind is LinkKind.Direct)
            {
                seenCall[target] = LinkKind.Direct;
            }
        }
    }

    private IEnumerable<ISymbol> FindImplementations(ISymbol member)
    {
        foreach (var type in types)
        {
            if (member.ContainingType.TypeKind is TypeKind.Interface)
            {
                if (type.FindImplementationForInterfaceMember(member) is { } implementation)
                {
                    yield return implementation;
                }
                continue;
            }

            foreach (var candidate in type.GetMembers(member.Name))
            {
                for (
                    var overridden = GetOverriddenMember(candidate);
                    overridden is not null;
                    overridden = GetOverriddenMember(overridden)
                )
                {
                    if (SymbolEqualityComparer.Default.Equals(overridden.OriginalDefinition, member.OriginalDefinition))
                    {
                        yield return candidate;
                        break;
                    }
                }
            }
        }
    }

    private static ISymbol? GetOverriddenMember(ISymbol member) =>
        member switch
        {
            IMethodSymbol method => method.OverriddenMethod,
            IPropertySymbol property => property.OverriddenProperty,
            _ => null,
        };

    private static bool IsDeclaredInSource(ISymbol symbol) =>
        symbol.DeclaringSyntaxReferences.Length > 0 && !symbol.IsImplicitlyDeclared;

    private static bool HasBody(IPropertySymbol property, CancellationToken cancellationToken) =>
        property.DeclaringSyntaxReferences.Any(reference =>
            reference.GetSyntax(cancellationToken) switch
            {
                PropertyDeclarationSyntax { ExpressionBody: not null } => true,
                IndexerDeclarationSyntax { ExpressionBody: not null } => true,
                BasePropertyDeclarationSyntax declaration => declaration.AccessorList?.Accessors.Any(accessor =>
                    accessor.Body is not null || accessor.ExpressionBody is not null
                )
                    is true,
                _ => false,
            }
        );
}