// Author: Simon.Gockner
// Created: 2025-12-01
// Copyright(c) 2025 SimonG. All Rights Reserved.

using System.Collections.Immutable;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace LightweightIocContainer.FactoryGenerator;

[Generator]
public class FactoryGenerator : IIncrementalGenerator
{
    private const string EXTENSION_CLASS_NAME = "FactoryExtensions";
    private const string BUILDER_CLASS_NAME = "FactoryBuilder";
    
    private const string GENERATED_FILE_HEADER = "//---GENERATED by FactoryGenerator! DO NOT EDIT!---";
    private const string INDENT = "    ";
    
    private const string CLEAR_MULTITON_INSTANCE_METHOD_NAME = "ClearMultitonInstance";
    
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        string? classNamespace = typeof(FactoryGenerator).Namespace;
        context.RegisterPostInitializationOutput(c => c.AddSource($"{EXTENSION_CLASS_NAME}.g.cs", GenerateFactoryExtensionsClass(classNamespace, EXTENSION_CLASS_NAME)));

        IncrementalValuesProvider<ITypeSymbol?> syntaxProvider = context.SyntaxProvider.CreateSyntaxProvider(IsCallToGenerateFactory, GetTypeArgument);
        
        context.RegisterSourceOutput(syntaxProvider.Collect(), GenerateTypeDependentClasses);
        context.RegisterSourceOutput(syntaxProvider.Collect(), GenerateFactory);
    }

    private string GenerateFactoryExtensionsClass(string? classNamespace, string className)
    {
        StringBuilder stringBuilder = new();
        
        stringBuilder.AppendLine(GENERATED_FILE_HEADER);
        stringBuilder.AppendLine();

        stringBuilder.AppendLine("using LightweightIocContainer.Interfaces.Factories;");
        stringBuilder.AppendLine("using LightweightIocContainer.Interfaces.Registrations;");
        stringBuilder.AppendLine();
        
        if (!string.IsNullOrEmpty(classNamespace))
        {
            stringBuilder.AppendLine($"namespace {classNamespace};");
            stringBuilder.AppendLine();
        }
        
        stringBuilder.AppendLine($"public static class {className}");
        stringBuilder.AppendLine("{");
        
        stringBuilder.AppendLine($"{INDENT}public static IRegistrationBase WithGeneratedFactory<TFactory>(this IRegistrationBase registration)");
        stringBuilder.AppendLine($"{INDENT}{{");
        stringBuilder.AppendLine($"{INDENT}{INDENT}FactoryBuilder factoryBuilder = new();");
        stringBuilder.AppendLine($"{INDENT}{INDENT}registration.AddGeneratedFactory<TFactory>(factoryBuilder);");
        stringBuilder.AppendLine();
        stringBuilder.AppendLine($"{INDENT}{INDENT}return registration;");
        stringBuilder.AppendLine($"{INDENT}}}");
        
        stringBuilder.AppendLine("}");
        
        return stringBuilder.ToString();
    }

    private bool IsCallToGenerateFactory(SyntaxNode node, CancellationToken cancellationToken)
    {
        if (!node.IsKind(SyntaxKind.GenericName) || node is not GenericNameSyntax genericNameSyntax)
            return false;

        if (!genericNameSyntax.ToString().StartsWith("WithGeneratedFactory"))
            return false;

        if (genericNameSyntax.TypeArgumentList.Arguments[0] is not (IdentifierNameSyntax or GenericNameSyntax))
            return false;

        return true;
    }
    
    private ITypeSymbol? GetTypeArgument(GeneratorSyntaxContext syntaxContext, CancellationToken cancellationToken)
    {
        if (syntaxContext.Node is not GenericNameSyntax genericNameSyntax)
            return null;
        
        if (genericNameSyntax.TypeArgumentList.Arguments[0] is not SimpleNameSyntax nameSyntax)
            return null;

        if (syntaxContext.SemanticModel.GetSymbolInfo(nameSyntax).Symbol is not ITypeSymbol typeSymbol)
            return null;

        return typeSymbol;
    }
    
    private void GenerateTypeDependentClasses(SourceProductionContext context, ImmutableArray<ITypeSymbol?> types)
    {
        string? classNamespace = typeof(FactoryGenerator).Namespace;
        context.AddSource($"{BUILDER_CLASS_NAME}.g.cs", GenerateBuilderClassSourceCode(classNamespace, types));
    }
    
    private void GenerateFactory(SourceProductionContext context, ImmutableArray<ITypeSymbol?> types)
    {
        foreach (ISymbol? symbol in types.Distinct(SymbolEqualityComparer.IncludeNullability))
        {
            if (symbol is not ITypeSymbol typeSymbol)
                continue;
        
            context.AddSource($"Generated{GetGenericFileName(typeSymbol)}.g.cs", GenerateFactorySourceCode(typeSymbol));
        }
    }

    private string GenerateBuilderClassSourceCode(string? classNamespace, ImmutableArray<ITypeSymbol?> types)
    {
        StringBuilder stringBuilder = new();

        stringBuilder.AppendLine(GENERATED_FILE_HEADER);
        stringBuilder.AppendLine();

        stringBuilder.AppendLine("using LightweightIocContainer.Interfaces.Factories;");
        
        foreach (string typeNamespace in GetNamespacesOfTypes(types)) 
            stringBuilder.AppendLine($"using {typeNamespace};");
        
        stringBuilder.AppendLine();

        if (classNamespace is not null)
        {
            stringBuilder.AppendLine($"namespace {classNamespace};");
            stringBuilder.AppendLine();
        }
        
        stringBuilder.AppendLine("public class FactoryBuilder : IFactoryBuilder");
        stringBuilder.AppendLine("{");
        stringBuilder.AppendLine($"{INDENT}public TFactory Create<TFactory>(IocContainer container)");
        stringBuilder.AppendLine($"{INDENT}{{");

        foreach (ISymbol? symbol in types.Distinct(SymbolEqualityComparer.IncludeNullability))
        {
            if (symbol is not ITypeSymbol type)
                continue;
            
            stringBuilder.AppendLine($"{INDENT}{INDENT}if (typeof(TFactory) == typeof({GetTypeText(type, false)}))");
            stringBuilder.AppendLine($"{INDENT}{INDENT}{{");
            stringBuilder.AppendLine($"{INDENT}{INDENT}{INDENT}return (TFactory) (object) new Generated{GetTypeText(type, false)}(container);");
            stringBuilder.AppendLine($"{INDENT}{INDENT}}}");
            stringBuilder.AppendLine();
        }
        
        stringBuilder.AppendLine($"{INDENT}{INDENT}throw new Exception(\"Invalid type.\");");
        
        stringBuilder.AppendLine($"{INDENT}}}");
        stringBuilder.AppendLine("}");

        return stringBuilder.ToString();
    }
    
    private string GenerateFactorySourceCode(ITypeSymbol typeSymbol)
    {
        string? typeNamespace = typeSymbol.ContainingNamespace.IsGlobalNamespace ? null : typeSymbol.ContainingNamespace.ToString();
        
        StringBuilder stringBuilder = new();

        stringBuilder.AppendLine(GENERATED_FILE_HEADER);
        stringBuilder.AppendLine();
        
        stringBuilder.AppendLine("#nullable enable");
        stringBuilder.AppendLine();

        stringBuilder.AppendLine("using LightweightIocContainer;");
        
        ImmutableArray<ISymbol> members = typeSymbol.GetMembers();

        List<string?> namespaces = [];
        foreach (ISymbol? member in members)
        {
            if (member is not IMethodSymbol method)
                continue;

            if (!method.ReturnsVoid) 
                namespaces.AddRange(GetNamespacesOfType(method.ReturnType));

            namespaces.AddRange(method.Parameters.SelectMany(p => GetNamespacesOfType(p.Type)));
        }

        foreach (string @namespace in namespaces.Distinct().OfType<string>().OrderBy(n => n)) 
            stringBuilder.AppendLine($"using {@namespace};");
        
        stringBuilder.AppendLine();

        if (typeNamespace is not null)
        {
            stringBuilder.AppendLine($"namespace {typeNamespace};");
            stringBuilder.AppendLine();
        }

        stringBuilder.Append($"public class Generated{GetGenericTypeText(typeSymbol)}(IocContainer container) : {GetGenericTypeText(typeSymbol)}");

        if (typeSymbol is INamedTypeSymbol { IsGenericType: true } genericTypeSymbol)
            stringBuilder.Append(GetParameterConstraintsText(genericTypeSymbol));

        stringBuilder.AppendLine();
        stringBuilder.AppendLine("{");

        foreach (ISymbol? member in members)
        {
            if (member is not IMethodSymbol method)
                continue;

            if (!method.ReturnsVoid) //create method
            {
                stringBuilder.Append($"{INDENT}public {GetTypeText(method.OriginalDefinition.ReturnType)} {method.Name}");
                
                if (method.IsGenericMethod) 
                    stringBuilder.Append(GetGenericParameters(method.OriginalDefinition));
                
                stringBuilder.Append($"({string.Join(", ", method.OriginalDefinition.Parameters.Select(GetParameterText))})");

                if (method.IsGenericMethod) 
                    stringBuilder.Append(GetParameterConstraintsText(method.OriginalDefinition));
                
                stringBuilder.AppendLine();
                stringBuilder.AppendLine($"{INDENT}{{");

                foreach (IParameterSymbol parameter in method.OriginalDefinition.Parameters)
                {
                    stringBuilder.AppendLine($"{INDENT}{INDENT}object? {parameter.Name}Value = {parameter.Name};");
                    stringBuilder.AppendLine($"{INDENT}{INDENT}if ({parameter.Name}Value is null)");
                    stringBuilder.AppendLine($"{INDENT}{INDENT}{INDENT}{parameter.Name}Value = new NullParameter(typeof({GetTypeText(parameter.Type, false)}));");
                    stringBuilder.AppendLine();
                }
                
                //don't use getTypeText here, because we need the raw type name for Task<>
                if (method.OriginalDefinition.ReturnType is INamedTypeSymbol { IsGenericType: true } namedTypeSymbol) 
                {
                    if (method.OriginalDefinition.ReturnType.Name == "Task") 
                        stringBuilder.Append($"{INDENT}{INDENT}return container.FactoryResolveAsync{GetGenericArguments(namedTypeSymbol)}(");
                    else
                        stringBuilder.Append($"{INDENT}{INDENT}return container.FactoryResolve<{method.OriginalDefinition.ReturnType.Name}{GetGenericArguments(namedTypeSymbol)}>(");
                }
                else
                    stringBuilder.Append($"{INDENT}{INDENT}return container.FactoryResolve<{method.OriginalDefinition.ReturnType.Name}>(");
                
                stringBuilder.Append(string.Join(", ", method.OriginalDefinition.Parameters.Select(p => $"{p.Name}Value")));
                stringBuilder.AppendLine(");");
                
                stringBuilder.AppendLine($"{INDENT}}}");
            }
            else if (method is { Name: CLEAR_MULTITON_INSTANCE_METHOD_NAME, IsGenericMethod: true })
            {
                stringBuilder.Append($"{INDENT}public void {method.Name}{GetGenericParameters(method)}()");
                stringBuilder.Append(GetParameterConstraintsText(method));
                
                stringBuilder.AppendLine($" => container.ClearMultitonInstances{GetGenericArguments(method)}();");
            }
            
            if (members.IndexOf(member) < members.Length - 1) //only append empty line if not the last member
                stringBuilder.AppendLine();
        }

        stringBuilder.AppendLine("}");
        
        return stringBuilder.ToString();
    }
    
    private string GetGenericFileName(ITypeSymbol typeSymbol)
    {
        StringBuilder stringBuilder = new();
        stringBuilder.Append(typeSymbol.Name);
        
        if (typeSymbol is INamedTypeSymbol { IsGenericType: true } namedTypeSymbol)
            stringBuilder.Append(string.Join("", namedTypeSymbol.TypeArguments.Select(t => CapitalizeFirstLetter(t.Name))));
        
        return stringBuilder.ToString();
    }

    private string? GetNamespaceOfType(ITypeSymbol typeSymbol) => typeSymbol.ContainingNamespace.IsGlobalNamespace ? null : typeSymbol.ContainingNamespace.ToString();
    private List<string?> GetNamespacesOfType(ITypeSymbol typeSymbol)
    {
        List<string?> namespaces = [];
        if (typeSymbol is INamedTypeSymbol { IsGenericType: true } namedTypeSymbol)
        {
            namespaces.AddRange(namedTypeSymbol.TypeArguments.SelectMany(GetNamespacesOfType));
                    
            if (typeSymbol.Name != "Task")
                namespaces.Add(GetNamespaceOfType(typeSymbol));
        }
        else
            namespaces.Add(GetNamespaceOfType(typeSymbol));

        return namespaces;
    }
    
    private IEnumerable<string> GetNamespacesOfTypes(ImmutableArray<ITypeSymbol?> types) =>
        types.OfType<ITypeSymbol>()
            .SelectMany(GetNamespacesOfType)
            .OfType<string>()
            .Distinct();

    private string GetTypeText(ITypeSymbol typeSymbol) => GetTypeText(typeSymbol, true);
    private string GetTypeText(ITypeSymbol typeSymbol, bool allowNullable)
    {
        StringBuilder stringBuilder = new();
        stringBuilder.Append(typeSymbol.Name);
        
        if (typeSymbol is INamedTypeSymbol { IsGenericType: true } namedTypeSymbol)
            stringBuilder.Append(GetGenericArguments(namedTypeSymbol));
        
        if (allowNullable && typeSymbol.NullableAnnotation == NullableAnnotation.Annotated)
            stringBuilder.Append("?");
        
        return stringBuilder.ToString();
    }

    private string GetGenericTypeText(ITypeSymbol typeSymbol)
    {
        StringBuilder stringBuilder = new();
        stringBuilder.Append(typeSymbol.Name);
        
        if (typeSymbol is INamedTypeSymbol { IsGenericType: true } namedTypeSymbol)
            stringBuilder.Append(GetGenericParameters(namedTypeSymbol));
        
        return stringBuilder.ToString();
    }

    private string GetParameterText(IParameterSymbol parameter) => $"{GetTypeText(parameter.Type)} {parameter.Name}";
    private string GetGenericArguments(INamedTypeSymbol namedTypeSymbol) => GetGenericArguments(namedTypeSymbol.TypeArguments);
    private string GetGenericArguments(IMethodSymbol methodSymbol) => GetGenericArguments(methodSymbol.TypeArguments);
    private string GetGenericArguments(ImmutableArray<ITypeSymbol> typeArguments) => $"<{string.Join(", ", typeArguments.Select(GetTypeText))}>";

    private string GetGenericParameters(INamedTypeSymbol namedTypeSymbol) => GetGenericParameters(namedTypeSymbol.TypeParameters);
    private string GetGenericParameters(IMethodSymbol methodSymbol) => GetGenericParameters(methodSymbol.TypeParameters);
    private string GetGenericParameters(ImmutableArray<ITypeParameterSymbol> typeParameters) => $"<{string.Join(", ", typeParameters.Select(GetTypeText))}>";

    private string GetParameterConstraintsText(INamedTypeSymbol namedTypeSymbol) => GetParameterConstraintsText(namedTypeSymbol.TypeParameters);
    private string GetParameterConstraintsText(IMethodSymbol method) => GetParameterConstraintsText(method.TypeParameters);
    private string GetParameterConstraintsText(ImmutableArray<ITypeParameterSymbol> typeParameters)
    {
        StringBuilder stringBuilder = new();
        
        foreach (ITypeParameterSymbol typeParameter in typeParameters)
        {
            List<string> parameterConstraints = GetParameterConstraints(typeParameter);
            if (parameterConstraints.Count == 0)
                continue;
                        
            stringBuilder.Append($" where {typeParameter.Name} : {string.Join(", ", parameterConstraints)}");
        }
        
        return stringBuilder.ToString();
    }

    private List<string> GetParameterConstraints(ITypeParameterSymbol typeParameterSymbol)
    {
        List<string> constraints = [];
        if (typeParameterSymbol.HasReferenceTypeConstraint)
            constraints.Add("class");
        
        if (typeParameterSymbol.HasValueTypeConstraint)
            constraints.Add("struct");
        
        if (typeParameterSymbol.HasUnmanagedTypeConstraint)
            constraints.Add("unmanaged");
        
        if (typeParameterSymbol.HasNotNullConstraint)
            constraints.Add("notnull");
        
        foreach (ITypeSymbol constraintType in typeParameterSymbol.ConstraintTypes)
        {
            if (constraintType is INamedTypeSymbol { IsGenericType: true } namedTypeSymbol)
                constraints.Add($"{constraintType.Name}{GetGenericArguments(namedTypeSymbol)}");
            else
                constraints.Add(constraintType.Name);
        }
        
        if (typeParameterSymbol.HasConstructorConstraint)
            constraints.Add("new()");
        
        return constraints;
    }
    
    private string CapitalizeFirstLetter(string input)
    {
        if (string.IsNullOrEmpty(input))
            return input;
        
        return char.ToUpper(input[0]) + input.Substring(1);
    }
}