ContextCodeGenerator.cs

using System.Collections.Generic;
using System.Text;
using System.Linq;
using Microsoft.CodeAnalysis;

namespace FastCloner.SourceGenerator;

internal sealed class ContextCodeGenerator
{
    private readonly ContextModel _model;
    private readonly StringBuilder _sb = new StringBuilder();
    private readonly Dictionary<string, string> _sharedMethodNames = new Dictionary<string, string>();
    private readonly HashSet<string> _sharedNeededHelpers = [];
    private Dictionary<string, bool> _needsState = new Dictionary<string, bool>();

    public ContextCodeGenerator(ContextModel model)
    {
        _model = model;
    }

    public string Generate()
    {
        AnalyzeCircularDependencies();

        _sb.AppendLine("// <auto-generated/>");
        _sb.AppendLine($"// Generated: {System.DateTime.UtcNow:yyyy-MM-dd HH:mm:ss} UTC");
        _sb.AppendLine("#nullable enable");
        _sb.AppendLine("#pragma warning disable");
        _sb.AppendLine();
        _sb.AppendLine("using System;");
        _sb.AppendLine("using System.Collections.Generic;");
        _sb.AppendLine("using System.Reflection;");
 
        if (_model.IsFastClonerAvailable)
        {
            _sb.AppendLine("using FastCloner;");
        }
        
        _sb.AppendLine();

        if (!string.IsNullOrEmpty(_model.Namespace))
        {
            _sb.AppendLine($"namespace {_model.Namespace}");
            _sb.AppendLine("{");
        }

        _sb.AppendLine($"    public partial class {_model.Name}");
        _sb.AppendLine("    {");

        // Generate methods for each registered type
        foreach (TypeModel? typeModel in _model.RegisteredTypes)
        {
            GenerateTypeCloner(typeModel);
        }

        // Generate dispatcher
        GenerateDispatcher();
        GenerateIsHandled();
        GenerateTryClone();

        // Generate State class (shared)
        GenerateStateClass();

        _sb.AppendLine("    }");

        if (!string.IsNullOrEmpty(_model.Namespace))
        {
            _sb.AppendLine("}");
        }

        return _sb.ToString();
    }
    
    private void GenerateTypeCloner(TypeModel model)
    {
        CloneGeneratorContext ctx = new CloneGeneratorContext(model, _sharedMethodNames, _sharedNeededHelpers);
        ctx.NeedsStateClass = true; // Use shared state class logic
        ctx.UseStaticMethods = false; // Use instance methods for context

        // Populate overrides from our analysis
        foreach (KeyValuePair<string, bool> kvp in _needsState)
        {
            ctx.SetCircularReferenceOverride(kvp.Key, kvp.Value);
        }

        // Pre-register all registered types to map to "Clone" method
        foreach (TypeModel? regType in _model.RegisteredTypes)
        {
             ctx.RegisterExternalMethod(regType.FullyQualifiedName, "Clone");
             ctx.RegisterImplicitType(regType);
        }

        string typeName = model.FullyQualifiedName;

        if (model.CircularAnalysisLog.Count > 0)
        {
             _sb.AppendLine("        // Circular Analysis Log:");
             foreach (string? line in model.CircularAnalysisLog)
             {
                 _sb.AppendLine($"        // {line}");
             }
        }

        // Determine if we really need state tracking for this type
        bool needsState = _needsState.ContainsKey(model.FullyQualifiedName) && _needsState[model.FullyQualifiedName];

        // Only generate NotNullIfNotNull attribute if it's available in the runtime (not from polyfill)
        string notNullAttr = CloneGeneratorContext.NotNullIfNotNullAttr(_model.CodeAnalysisAvailable);
        if (!string.IsNullOrEmpty(notNullAttr))
            _sb.AppendLine($"        {notNullAttr}");
        _sb.AppendLine($"        public {typeName}? Clone({typeName}? source)");
        _sb.AppendLine("        {");
        
        if (!needsState)
        {
             _sb.AppendLine("            if (source == null) return null;");
             GenerateCloneBody(ctx, typeName, false);
             _sb.Append(ctx.Source.ToString());
             ctx.Source.Clear();
        }
        else
        {
             // Fix: if type needs state, the public method MUST initialize it
             _sb.AppendLine($"            return Clone(source, new FcGeneratedCloneState());");
        }
        _sb.AppendLine("        }");
        _sb.AppendLine();

        if (needsState)
        {
             _sb.AppendLine($"        private {typeName}? Clone({typeName}? source, FcGeneratedCloneState state)");
             _sb.AppendLine("        {");
             _sb.AppendLine("            if (source == null) return null;");
             // Fix: Remove the state ?? new check - state is always passed and managed by caller
             
             if (!model.IsStruct)
             {
                 _sb.AppendLine("            var known = state.GetKnownRef(source);");
                 _sb.AppendLine($"            if (known != null) return ({typeName})known;");
             }

             GenerateCloneBody(ctx, typeName, true, "state");
             _sb.Append(ctx.Source.ToString());
             ctx.Source.Clear();
             _sb.AppendLine("        }");
             _sb.AppendLine();
        }
        
        // Generate helpers that were triggered by this type
        CollectionHelperGenerator.GenerateHelpers(ctx);
        
        // Append all generated code from context (helpers)
        _sb.Append(ctx.Source.ToString());
    }
    
    private void GenerateCloneBody(CloneGeneratorContext ctx, string typeName, bool useState, string? stateVarName = null)
    {
         if (ctx.Model.IsStruct)
         {
             WriteStructCloneBody(ctx, stateVarName ?? "state");
         }
         else
         {
             WriteClassCloneBody(ctx, typeName, useState, stateVarName);
         }
    }

    private void WriteStructCloneBody(CloneGeneratorContext ctx, string stateVarName)
    {
        StringBuilder sb = ctx.Source;
        sb.AppendLine($"            var result = source;");
        sb.AppendLine();

        foreach (MemberModel member in ctx.Model.Members)
        {
            MemberCloneGenerator.WriteMemberCloning(ctx, member, "result", "source", stateVarName);
        }

        sb.AppendLine("            return result;");
    }

    private void WriteClassCloneBody(CloneGeneratorContext ctx, string typeName, bool useState, string? stateVarName)
    {
        ClassCloneBodyGenerator.WriteClassCloneBody(ctx, typeName, useState, stateVarName, useNullConditional: false);
    }

    private void GenerateDispatcher()
    {
        _sb.AppendLine("        public override object Clone(object input)");
        _sb.AppendLine("        {");
        foreach (TypeModel? type in _model.RegisteredTypes)
        {
            // Call public Clone method which handles state initialization
            _sb.AppendLine($"            if (input is {type.FullyQualifiedName} t{type.Name}) return Clone(t{type.Name})!;");
        }
        
        // Fallback to reflection-based FastCloner if available, otherwise throw
        if (_model.IsFastClonerAvailable)
        {
            _sb.AppendLine($"            return {CloneGeneratorContext.FastClonerDeepCloneCall("input")}!;");
        }
        else
        {
            _sb.AppendLine("            throw new System.NotSupportedException(\"Cannot clone unregistered type without FastCloner runtime library.\");");
        }
        
        _sb.AppendLine("        }");
        _sb.AppendLine();
    }

    private void GenerateIsHandled()
    {
        _sb.AppendLine("        public override bool IsHandled(Type type)");
        _sb.AppendLine("        {");
        foreach (TypeModel? type in _model.RegisteredTypes)
        {
            _sb.AppendLine($"            if (type == typeof({type.FullyQualifiedName})) return true;");
        }
        _sb.AppendLine("            return false;");
        _sb.AppendLine("        }");
        _sb.AppendLine();
    }

    private void GenerateTryClone()
    {
        // Only generate NotNullWhen attribute if NotNullIfNotNull is available (they're in the same namespace)
        if (_model.CodeAnalysisAvailable)
        {
            _sb.AppendLine("        public override bool TryClone(object input, [global::System.Diagnostics.CodeAnalysis.NotNullWhen(true)] out object? clone)");
        }
        else
        {
            _sb.AppendLine("        public override bool TryClone(object input, out object? clone)");
        }
        _sb.AppendLine("        {");
        _sb.AppendLine("            if (input == null)");
        _sb.AppendLine("            {");
        _sb.AppendLine("                clone = null;");
        _sb.AppendLine("                return false;");
        _sb.AppendLine("            }");
        _sb.AppendLine();
        foreach (TypeModel? type in _model.RegisteredTypes)
        {
            _sb.AppendLine($"            if (input is {type.FullyQualifiedName} t{type.Name})");
            _sb.AppendLine("            {");
            _sb.AppendLine($"                clone = Clone(t{type.Name});");
            _sb.AppendLine("                return true;");
            _sb.AppendLine("            }");
        }
        _sb.AppendLine("            clone = null;");
        _sb.AppendLine("            return false;");
        _sb.AppendLine("        }");
        _sb.AppendLine();
    }

    private void GenerateStateClass()
    {
        // Only generate state class if at least one type needs circular reference tracking
        if (!_needsState.Any(kvp => kvp.Value))
            return;

        _sb.AppendLine("        /// <summary>");
        _sb.AppendLine("        /// State for tracking circular references during cloning.");
        _sb.AppendLine("        /// </summary>");
        _sb.AppendLine("        private class FcGeneratedCloneState");
        _sb.AppendLine("        {");
        _sb.AppendLine("            private readonly Dictionary<object, object> _knownRefs = new Dictionary<object, object>(new ReferenceEqualityComparer());");
        _sb.AppendLine();
        _sb.AppendLine("            public void AddKnownRef(object original, object clone)");
        _sb.AppendLine("            {");
        _sb.AppendLine("                if (original != null)");
        _sb.AppendLine("                {");
        _sb.AppendLine("                    _knownRefs[original] = clone;");
        _sb.AppendLine("                }");
        _sb.AppendLine("            }");
        _sb.AppendLine();
        _sb.AppendLine("            public object? GetKnownRef(object original)");
        _sb.AppendLine("            {");
        _sb.AppendLine("                if (original == null) return null;");
        _sb.AppendLine("                return _knownRefs.TryGetValue(original, out var clone) ? clone : null;");
        _sb.AppendLine("            }");
        _sb.AppendLine();
        _sb.AppendLine("            private sealed class ReferenceEqualityComparer : IEqualityComparer<object>");
        _sb.AppendLine("            {");
        _sb.AppendLine("                bool IEqualityComparer<object>.Equals(object? x, object? y) => ReferenceEquals(x, y);");
        _sb.AppendLine("                int IEqualityComparer<object>.GetHashCode(object obj) => System.Runtime.CompilerServices.RuntimeHelpers.GetHashCode(obj);");
        _sb.AppendLine("            }");
        _sb.AppendLine("        }");
    }

    private void AnalyzeCircularDependencies()
    {
        // 1. Build Graph
        Dictionary<string, HashSet<string>> adjacency = new Dictionary<string, HashSet<string>>();
        Dictionary<string, TypeModel> registeredTypes = _model.RegisteredTypes.ToDictionary(t => t.FullyQualifiedName, t => t);

        foreach (TypeModel? type in _model.RegisteredTypes)
        {
            HashSet<string> deps = [];
            adjacency[type.FullyQualifiedName] = deps;

            foreach (MemberModel member in type.Members)
            {
                AddDependency(deps, member.TypeFullName, registeredTypes);
                AddDependency(deps, member.ElementTypeName, registeredTypes);
                AddDependency(deps, member.KeyTypeName, registeredTypes);
                AddDependency(deps, member.ValueTypeName, registeredTypes);
            }
        }

        // 2. Find nodes that reach a cycle
        Dictionary<string, VisitStatus> status = new Dictionary<string, VisitStatus>();
        Dictionary<string, bool> reachesCycle = new Dictionary<string, bool>();

        foreach (TypeModel? type in _model.RegisteredTypes)
        {
            if (!status.ContainsKey(type.FullyQualifiedName))
            {
                CheckCycleDFS(type.FullyQualifiedName, adjacency, status, reachesCycle);
            }
        }

        _needsState = reachesCycle;
    }

    private void AddDependency(HashSet<string> deps, string? typeName, Dictionary<string, TypeModel> registered)
    {
        if (!string.IsNullOrEmpty(typeName) && registered.ContainsKey(typeName!))
        {
            deps.Add(typeName!);
        }
    }

    private enum VisitStatus { Unvisited, Visiting, Visited }

    private bool CheckCycleDFS(string u, Dictionary<string, HashSet<string>> adj, Dictionary<string, VisitStatus> status, Dictionary<string, bool> reachesCycle)
    {
        status[u] = VisitStatus.Visiting;
        bool cycleFound = false;

        if (adj.TryGetValue(u, out HashSet<string>? neighbors))
        {
            foreach (string? v in neighbors)
            {
                if (!status.ContainsKey(v)) status[v] = VisitStatus.Unvisited;

                if (status[v] == VisitStatus.Visiting)
                {
                    // Back edge found -> Cycle
                    cycleFound = true;
                }
                else if (status[v] == VisitStatus.Visited)
                {
                    if (reachesCycle.ContainsKey(v) && reachesCycle[v])
                    {
                        cycleFound = true;
                    }
                }
                else // Unvisited
                {
                    if (CheckCycleDFS(v, adj, status, reachesCycle))
                    {
                        cycleFound = true;
                    }
                }
            }
        }

        status[u] = VisitStatus.Visited;
        reachesCycle[u] = cycleFound;
        return cycleFound;
    }
}