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

namespace Arch.System.SourceGenerator;

public static class QueryUtils
{

    /// <summary>
    ///     Appends the first elements of the types specified in the <see cref="parameterSymbols"/> from the previous specified arrays.
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="parameterSymbols">The <see cref="IEnumerable{T}"/> list of <see cref="IParameterSymbol"/>s which we wanna append the first elements for.</param>
    /// <returns></returns>
    public static StringBuilder GetFirstElements(this StringBuilder sb, IEnumerable<IParameterSymbol> parameterSymbols)
    {
      
        foreach (var symbol in parameterSymbols)
            if(symbol.Type.Name is not "Entity" || !symbol.GetAttributes().Any(data => data.AttributeClass.Name.Contains("Data"))) // Prevent entity being added to the type array
                sb.AppendLine($"ref var {symbol.Type.Name.ToLower()}FirstElement = ref chunk.GetFirst<{symbol.Type}>();");

        return sb;
    }
    
    /// <summary>
    ///     Appends the components of the types specified in the <see cref="parameterSymbols"/> from the previous specified first elements.
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="parameterSymbols">The <see cref="IEnumerable{T}"/> list of <see cref="IParameterSymbol"/>s which we wanna append the components for.</param>
    /// <returns></returns>
    public static StringBuilder GetComponents(this StringBuilder sb, IEnumerable<IParameterSymbol> parameterSymbols)
    {
        foreach (var symbol in parameterSymbols)
            if(symbol.Type.Name is not "Entity") // Prevent entity being added to the type array
                sb.AppendLine($"ref var {symbol.Name.ToLower()} = ref Unsafe.Add(ref {symbol.Type.Name.ToLower()}FirstElement, entityIndex);");

        return sb;
    }
    
    /// <summary>
    ///     Inserts the types defined in the <see cref="parameterSymbols"/> as parameters in a method.
    ///     <example>ref position, out velocity,...</example>
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="parameterSymbols">The <see cref="IEnumerable{T}"/> of <see cref="IParameterSymbol"/>s which we wanna insert.</param>
    /// <returns></returns>
    public static StringBuilder InsertParams(this StringBuilder sb, IEnumerable<IParameterSymbol> parameterSymbols)
    {
        foreach (var symbol in parameterSymbols)
            sb.Append($"{CommonUtils.RefKindToString(symbol.RefKind)} {symbol.Name.ToLower()},");
        
        if(sb.Length > 0) sb.Length--;
        return sb;
    }
    
    /// <summary>
    ///     Creates a ComponentType array from the <see cref="parameterSymbols"/> passed through.
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/>.</param>
    /// <param name="parameterSymbols">The <see cref="IList{T}"/> with <see cref="ITypeSymbol"/>s which we wanna create a ComponentType array for.</param>
    /// <returns></returns>
    public static StringBuilder GetTypeArray(this StringBuilder sb, IList<ITypeSymbol> parameterSymbols)
    {
        if (parameterSymbols.Count == 0)
        {
            sb.Append("Array.Empty<ComponentType>()");
            return sb;
        }

        sb.Append("new ComponentType[]{");
        
        foreach (var symbol in parameterSymbols)
            if(symbol.Name is not "Entity") // Prevent entity being added to the type array
                sb.Append($"typeof({symbol}),");
        
        if (sb.Length > 0) sb.Length -= 1;
        sb.Append('}');

        return sb;
    }


    /// <summary>
    ///     Appends a set of <see cref="parameterSymbols"/> if they are marked by the data attribute.
    ///     <example>ref gameTime, out somePassedList,...</example>
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="parameterSymbols">The <see cref="IEnumerable{T}"/> of <see cref="IParameterSymbol"/>s which will be appended if they are marked with data.</param>
    /// <returns></returns>
    public static StringBuilder DataParameters(this StringBuilder sb, IEnumerable<IParameterSymbol> parameterSymbols)
    {
        sb.Append(',');
        foreach (var parameter in parameterSymbols)
        {
            if (parameter.GetAttributes().Any(attributeData => attributeData.AttributeClass.Name.Contains("Data")))
                sb.Append($"{CommonUtils.RefKindToString(parameter.RefKind)} {parameter.Type} {parameter.Name.ToLower()},");
        }
        sb.Length--;
        return sb;
    }
    
    /// <summary>
    ///     Appends method calls made with their important data parameters.
    ///     <example>someQuery(World, gameTime); ...</example>
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="methodNames">The <see cref="IEnumerable{T}"/> of methods which we wanna call.</param>
    /// <returns></returns>
    public static StringBuilder CallMethods(this StringBuilder sb, IEnumerable<IMethodSymbol> methodNames)
    {
        foreach (var method in methodNames)
        {
            var data = new StringBuilder();
            data.Append(',');
            foreach (var parameter in method.Parameters)
            {
                if (!parameter.GetAttributes().Any(attributeData => attributeData.AttributeClass.Name.Contains("Data"))) continue;
                data.Append($"{CommonUtils.RefKindToString(parameter.RefKind)} Data,");
                break;
            }
            data.Length--;
            sb.AppendLine($"{method.Name}Query(World {data});");   
        }
        return sb;
    }

    /// <summary>
    ///     Gets all the types of a <see cref="AttributeData"/> as <see cref="ITypeSymbol"/>s and adds them to a list.
    ///     If the attribute is generic it will add the generic parameters, if its non generic it will add the non generic types from the constructor.
    /// </summary>
    /// <param name="data">The <see cref="AttributeData"/>.</param>
    /// <param name="array">The <see cref="List{T}"/> where the found <see cref="ITypeSymbol"/>s are added to.</param>
    public static void GetAttributeTypes(AttributeData data, List<ITypeSymbol> array)
    {
        if (data is not null && data.AttributeClass.IsGenericType)
        {
            array.AddRange(data.AttributeClass.TypeArguments);
        }
        else if (data is not null && !data.AttributeClass.IsGenericType)
        {
            var constructorArguments = data.ConstructorArguments[0].Values;
            var constructorArgumentsTypes = constructorArguments.Select(constant => constant.Value as ITypeSymbol).ToList();
            array.AddRange(constructorArgumentsTypes);
        }
    }

    /// <summary>
    ///     Adds a query with an entity for a given annotated method. The attributes of these methods are used to generate the query.
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="methodSymbol">The <see cref="IMethodSymbol"/> which is annotated for source generation.</param>
    /// <returns></returns>
    public static StringBuilder AppendQueryMethod(this StringBuilder sb, IMethodSymbol methodSymbol)
    {

        // Check for entity param
        var entity = methodSymbol.Parameters.Any(symbol => symbol.Type.Name.Equals("Entity"));
        var entityParam = entity ? methodSymbol.Parameters.First(symbol => symbol.Type.Name.Equals("Entity")) : null;

        // Get attributes
        var attributeData = methodSymbol.GetAttributeData("All");
        var anyAttributeData = methodSymbol.GetAttributeData("Any");
        var noneAttributeData = methodSymbol.GetAttributeData("None");
        var exclusiveAttributeData = methodSymbol.GetAttributeData("Exclusive");
        
        // Get params / components except those marked with data or entities. 
        var components = methodSymbol.Parameters.ToList();
        components.RemoveAll(symbol => symbol.Type.Name.Equals("Entity"));                                                // Remove entitys 
        components.RemoveAll(symbol => symbol.GetAttributes().Any(data => data.AttributeClass.Name.Contains("Data")));    // Remove data annotated params
        
        // Create all query array
        var allArray = components.Select(symbol => symbol.Type).ToList();
        var anyArray = new List<ITypeSymbol>();
        var noneArray = new List<ITypeSymbol>();
        var exclusiveArray = new List<ITypeSymbol>();

        // Get All<...> or All(...) passed types and pass them to the arrays 
        GetAttributeTypes(attributeData, allArray);
        GetAttributeTypes(anyAttributeData, anyArray);
        GetAttributeTypes(noneAttributeData, noneArray);
        GetAttributeTypes(exclusiveAttributeData, exclusiveArray);
        
        // Remove doubles and entities 
        allArray = allArray.Distinct().ToList();
        anyArray = anyArray.Distinct().ToList();
        noneArray = noneArray.Distinct().ToList();
        exclusiveArray = exclusiveArray.Distinct().ToList();
        
        allArray.RemoveAll(symbol => symbol.Name.Equals("Entity")); 
        anyArray.RemoveAll(symbol => symbol.Name.Equals("Entity"));
        noneArray.RemoveAll(symbol => symbol.Name.Equals("Entity"));
        exclusiveArray.RemoveAll(symbol => symbol.Name.Equals("Entity"));

        // Create data modell and generate it
        var className = methodSymbol.ContainingSymbol.ToString();
        var queryMethod = new QueryMethod
        {
            IsGlobalNamespace = methodSymbol.ContainingNamespace.IsGlobalNamespace,
            Namespace = methodSymbol.ContainingNamespace.ToString(),
            ClassName = className.Substring(className.LastIndexOf('.')+1),
            
            IsStatic = methodSymbol.IsStatic,
            IsEntityQuery = entity,
            MethodName = methodSymbol.Name,
            
            EntityParameter = entityParam,
            Parameters = methodSymbol.Parameters,
            Components = components,
            
            AllFilteredTypes = allArray,
            AnyFilteredTypes = anyArray,
            NoneFilteredTypes = noneArray,
            ExclusiveFilteredTypes = exclusiveArray
        };
        
        return sb.AppendQueryMethod(ref queryMethod);
    }

    /// <summary>
    ///     Adds a query with an entity for a given annotated method. The attributes of these methods are used to generate the query.
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="queryMethod">The <see cref="QueryMethod"/> which is generated.</param>
    /// <returns></returns>
    public static StringBuilder AppendQueryMethod(this StringBuilder sb, ref QueryMethod queryMethod)
    {
        var staticModifier = queryMethod.IsStatic ? "static" : "";
        
        // Generate code 
        var data = new StringBuilder().DataParameters(queryMethod.Parameters);
        var getFirstElements = new StringBuilder().GetFirstElements(queryMethod.Components);
        var getComponents = new StringBuilder().GetComponents(queryMethod.Components);
        var insertParams = new StringBuilder().InsertParams(queryMethod.Parameters);
        
        var allTypeArray = new StringBuilder().GetTypeArray(queryMethod.AllFilteredTypes);
        var anyTypeArray = new StringBuilder().GetTypeArray(queryMethod.AnyFilteredTypes);
        var noneTypeArray = new StringBuilder().GetTypeArray(queryMethod.NoneFilteredTypes);
        var exclusiveTypeArray = new StringBuilder().GetTypeArray(queryMethod.ExclusiveFilteredTypes);

        var template = 
            $$"""
            using System;
            using System.Runtime.CompilerServices;
            using System.Runtime.InteropServices;
            using Arch.Core;
            using Arch.Core.Extensions;
            using Arch.Core.Utils;
            using ArrayExtensions = CommunityToolkit.HighPerformance.ArrayExtensions;
            using Component = Arch.Core.Utils.Component;
            {{(!queryMethod.IsGlobalNamespace ? $"namespace {queryMethod.Namespace} {{" : "")}}
                public {{staticModifier}} partial class {{queryMethod.ClassName}}{
                    
                    private {{staticModifier}} QueryDescription {{queryMethod.MethodName}}_QueryDescription = new QueryDescription{
                        All = {{allTypeArray}},
                        Any = {{anyTypeArray}},
                        None = {{noneTypeArray}},
                        Exclusive = {{exclusiveTypeArray}}
                    };

                    private {{staticModifier}} bool _{{queryMethod.MethodName}}_Initialized;
                    private {{staticModifier}} Query _{{queryMethod.MethodName}}_Query;

                    [MethodImpl(MethodImplOptions.AggressiveInlining)]
                    public {{staticModifier}} void {{queryMethod.MethodName}}Query(World world {{data}}){
                     
                        if(!_{{queryMethod.MethodName}}_Initialized){
                            _{{queryMethod.MethodName}}_Query = world.Query(in {{queryMethod.MethodName}}_QueryDescription);
                            _{{queryMethod.MethodName}}_Initialized = true;
                        }

                        foreach(ref var chunk in _{{queryMethod.MethodName}}_Query.GetChunkIterator()){
                            
                            var chunkSize = chunk.Size;
                            {{(queryMethod.IsEntityQuery ? "ref var entityFirstElement = ref chunk.Entity(0);" : "")}}
                            {{getFirstElements}}

                            foreach(var entityIndex in chunk)
                            {
                                {{(queryMethod.IsEntityQuery ? $"ref readonly var {queryMethod.EntityParameter.Name.ToLower()} = ref Unsafe.Add(ref entityFirstElement, entityIndex);" : "")}}
                                {{getComponents}}
                                {{queryMethod.MethodName}}({{insertParams}});
                            }
                        }
                    }
                }
            {{(!queryMethod.IsGlobalNamespace ? "}" : "")}}
            """;

        sb.Append(template);
        return sb;
    }

    /// <summary>
    ///     Adds a basesystem that calls a bunch of query methods. 
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="classToMethod">The <see cref="KeyValuePair{TKey,TValue}"/> which maps all query methods to a common class containing them.</param>
    /// <returns></returns>
    public static StringBuilder AppendBaseSystem(this StringBuilder sb, KeyValuePair<ISymbol, List<IMethodSymbol>> classToMethod)
    {
        // Get BaseSystem class
        var classSymbol = classToMethod.Key as INamedTypeSymbol;

        INamedTypeSymbol? parentSymbol = null;
        var implementsUpdate = false;
        var type = classSymbol;
        while (type != null)
        {
            // Update was implemented by user, no need to do that by source generator.
            if (type.MemberNames.Contains("Update"))
                implementsUpdate = true;

            type = type.BaseType;

            // Ignore classes which do not derive from BaseSystem
            if (type?.Name == "BaseSystem")
            {
                parentSymbol = type;
                break;
            }
        }

        if (parentSymbol == null || implementsUpdate)
            return sb;

        // Get generic of BaseSystem
        var typeSymbol = parentSymbol.TypeArguments[1];

        // Generate basesystem.
        var baseSystem = new BaseSystem
        {
            Namespace = classSymbol.ContainingNamespace != null ? classSymbol.ContainingNamespace.ToString() : string.Empty,
            GenericType = typeSymbol,
            GenericTypeNamespace = typeSymbol.ContainingNamespace.ToString(),
            Name = classSymbol.Name,
            QueryMethods = classToMethod.Value,
        };
        return sb.AppendBaseSystem(ref baseSystem);
    }
    
    /// <summary>
    ///     Adds a basesystem that calls a bunch of query methods. 
    /// </summary>
    /// <param name="sb">The <see cref="StringBuilder"/> instance.</param>
    /// <param name="baseSystem">The <see cref="BaseSystem"/> which is generated.</param>
    /// <returns></returns>
    public static StringBuilder AppendBaseSystem(this StringBuilder sb, ref BaseSystem baseSystem)
    {
        var methodCalls = new StringBuilder().CallMethods(baseSystem.QueryMethods);
        var template =
            $$"""
            using System.Runtime.CompilerServices;
            using System.Runtime.InteropServices;
            using {{baseSystem.GenericTypeNamespace}};
            {{(baseSystem.Namespace != string.Empty ? $"namespace {baseSystem.Namespace} {{" : "")}}
                public partial class {{baseSystem.Name}}{
                        
                    [MethodImpl(MethodImplOptions.AggressiveInlining)]
                    public override void Update(in {{baseSystem.GenericType.ToDisplayString()}} {{baseSystem.GenericType.Name.ToLower()}}){
                        {{methodCalls}}
                    }
                }
            {{(baseSystem.Namespace != string.Empty ? "}" : "")}}
            """;
        return sb.Append(template);
    }
}