#if UNITY_EDITOR

using UnityEngine;
using UnityEditor;
using System;
using System.Collections.Generic;
using System.IO;
using System.Text;

public class PrefabClassGenerator
{
    // If OUTPUT_FOLDER is empty, generated scripts will be placed next to their prefab.
    //static string OUTPUT_FOLDER = "Assets/PrefabClasses/";
    static string OUTPUT_FOLDER = "";

    [MenuItem("Assets/Generate Prefab Class", validate = true)]
    private static bool ValidateGeneratePrefabClass()
    {
        var prefabs = FindAllPrefabsInSelection();

        return prefabs.Count > 0;
    }

    [MenuItem("Assets/Generate Prefab Class", false, 1)]
    private static void GeneratePrefabClass()
    {
        var prefabs = FindAllPrefabsInSelection();

        foreach(var prefab in prefabs)
        {
            CreatePrefabClass(prefab);
        }

        AssetDatabase.Refresh();
    }

    private static List<GameObject> FindAllPrefabsInSelection()
    {
        List<GameObject> prefabs = new List<GameObject>();

        if(Selection.activeObject != null)
        {
            var folderPath = AssetDatabase.GetAssetPath(Selection.activeObject);
            
            //if a folder is selected, find all prefabs in the folder recursively
            if(AssetDatabase.IsValidFolder(folderPath))
            {
                string[] assetGUIDs = AssetDatabase.FindAssets("", new[] { folderPath });

                foreach (string guid in assetGUIDs)
                {
                    string assetPath = AssetDatabase.GUIDToAssetPath(guid);
                    var go = AssetDatabase.LoadAssetAtPath<GameObject>(assetPath);

                    if(go != null) 
                    {
                        prefabs.Add(go);
                    } 
                }
            }
            else //find all the prefabs in our current selected objects
            {
                foreach(var go in Selection.gameObjects)
                {
                    if(PrefabUtility.IsPartOfPrefabAsset(go))
                    {
                        prefabs.Add(go); 
                    }
                }
            }
        }

        

        return prefabs;
    }

    [MenuItem("Tools/Generate Classes For All Prefabs")]
    public static void GenerateClassesForPrefabs()
    {
        if (!string.IsNullOrEmpty(OUTPUT_FOLDER) && !Directory.Exists(OUTPUT_FOLDER))
        {
            Directory.CreateDirectory(OUTPUT_FOLDER);
            AssetDatabase.Refresh();
        }

        string[] guids = AssetDatabase.FindAssets("t:Prefab", new[] { "Assets" });
        Debug.Log($"Found {guids.Length} prefabs.");

        foreach (string guid in guids)
        {
            string prefabPath = AssetDatabase.GUIDToAssetPath(guid);
            GameObject prefab = AssetDatabase.LoadAssetAtPath<GameObject>(prefabPath);
            if (prefab == null)
            {
                Debug.LogWarning($"Could not load prefab at path: {prefabPath}");
                continue;
            }
            CreatePrefabClass(prefab);
        }
        AssetDatabase.Refresh();
    }

    static void CreatePrefabClass(GameObject prefab)
    {
        string className = SanitizeClassName(prefab.name) + "Prefab";

        string folderPath = OUTPUT_FOLDER;
        if (string.IsNullOrEmpty(OUTPUT_FOLDER))
        {
            string prefabAssetPath = AssetDatabase.GetAssetPath(prefab);
            folderPath = Path.GetDirectoryName(prefabAssetPath);
        }
        string filePath = Path.Combine(folderPath, className + ".cs");

        string code = GenerateCodeForPrefab(prefab, className);
        File.WriteAllText(filePath, code);
        Debug.Log($"Generated {filePath}");
    }

    private static string GenerateCodeForPrefab(GameObject prefab, string className)
    {
        StringBuilder sb = new StringBuilder();
        sb.AppendLine("using UnityEngine;");
        sb.AppendLine("using System.Collections.Generic;");
        sb.AppendLine();
        sb.AppendLine($"public class {className} : MonoBehaviour");
        sb.AppendLine("{");

        HashSet<string> usedNames = new HashSet<string>();

        // Built-in component types that conflict with MonoBehaviour properties.
        HashSet<string> builtIn = new HashSet<string>()
        {
            "Rigidbody", "Rigidbody2D", "Camera", "Animation", "ConstantForce", "Collider",
            "Collider2D", "HingeJoint", "NetworkView", "GUIText", "GUITexture", "Audio"
        };

        List<string> propertyDeclarations = new List<string>();
        List<string> setupLines = new List<string>();

        // Process root components (excluding Transform)
        Component[] comps = prefab.GetComponents<Component>();
        foreach (var comp in comps)
        {
            if (comp == null || comp is Transform)
                continue;

            Type compType = comp.GetType();
            if (compType.Name == prefab.name || compType.Name == prefab.name + "Prefab")
                continue;

            string typeString = GetTypeString(compType);
            string shortTypeName = compType.Name;

            string propName = char.ToLowerInvariant(shortTypeName[0]) + shortTypeName.Substring(1);
            string baseName = propName;
            int duplicateCount = 1;
            while (usedNames.Contains(propName))
            {
                propName = baseName + duplicateCount;
                duplicateCount++;
            }
            usedNames.Add(propName);

            string modifier = builtIn.Contains(shortTypeName) ? "new " : "";
            propertyDeclarations.Add($"    {modifier}public {typeString} {propName} {{ get; protected set; }}");
            setupLines.Add($"        {propName} = GetComponent<{typeString}>();");
        }

        // Process immediate child transforms and group them by name.
        Dictionary<string, List<ChildInfo>> childGroups = new Dictionary<string, List<ChildInfo>>();
        CollectDirectChildGroups(prefab.transform, childGroups);

        // Build group info: for a single child, generate a single property;
        // for multiple children with the same name, generate a List<> property.
        var childGroupList = new List<(string groupName, string propertyName, string typeString, bool isList)>();
        foreach (var kvp in childGroups)
        {
            string groupName = kvp.Key;
            List<ChildInfo> children = kvp.Value;
            bool isList = children.Count > 1;
            string typeString = children[0].TypeName;
            string propName = SanitizePropertyName(groupName, usedNames);
            if (isList)
                propName += "List";
            usedNames.Add(propName);
            childGroupList.Add((groupName, propName, typeString, isList));

            if (isList)
                propertyDeclarations.Add($"    public List<{typeString}> {propName} {{ get; protected set; }} = new List<{typeString}>();");
            else
                propertyDeclarations.Add($"    public {typeString} {propName} {{ get; protected set; }}");
        }

        // Generate Setup() method with a switch statement for child transforms.

        if(childGroupList.Count > 0)
        {
            if(setupLines.Count > 0) setupLines.Add(""); //add spacing if there are other component lines above this

            setupLines.Add("        for (int i = 0; i < transform.childCount; i++)");
            setupLines.Add("        {");
            setupLines.Add("            var child = transform.GetChild(i);");
            setupLines.Add("            switch(child.name)");
            setupLines.Add("            {");

            foreach (var group in childGroupList)
            {
                setupLines.Add($"                case \"{group.groupName}\":");
                if (group.isList)
                {
                    if (group.typeString == GetTypeString(typeof(Transform)))
                        setupLines.Add($"                    {group.propertyName}.Add(child);");
                    else
                        setupLines.Add($"                    {group.propertyName}.Add(child.GetComponent<{group.typeString}>());");
                }
                else
                {
                    if (group.typeString == GetTypeString(typeof(Transform)))
                        setupLines.Add($"                    {group.propertyName} = child;");
                    else
                        setupLines.Add($"                    {group.propertyName} = child.GetComponent<{group.typeString}>();");
                }
                setupLines.Add("                    break;");
            }

            setupLines.Add("            }");
            setupLines.Add("        }");
        }

        foreach (string line in propertyDeclarations)
            sb.AppendLine(line);

        sb.AppendLine();
        sb.AppendLine("    virtual protected void Setup()");
        sb.AppendLine("    {");
        foreach (string line in setupLines)
            sb.AppendLine(line);
        sb.AppendLine("    }");
        sb.AppendLine();
        sb.AppendLine("    virtual protected void Awake()");
        sb.AppendLine("    {");
        sb.AppendLine("        Setup();");
        sb.AppendLine("    }");
        sb.AppendLine("}");

        return sb.ToString();
    }

    // Collects only the immediate children and groups them by name.
    private static void CollectDirectChildGroups(Transform parent, Dictionary<string, List<ChildInfo>> childGroups)
    {
        foreach (Transform child in parent)
        {
            string baseName = child.name.Replace(" ", "");
            if (string.IsNullOrEmpty(baseName))
                continue;

            string typeString = GetTypeString(typeof(Transform));
            Component[] comps = child.GetComponents<Component>();
            foreach (var comp in comps)
            {
                if (comp == null || comp is Transform)
                    continue;
                typeString = GetTypeString(comp.GetType());
                break;
            }

            ChildInfo info = new ChildInfo(child.name, typeString);
            if (!childGroups.ContainsKey(child.name))
                childGroups[child.name] = new List<ChildInfo>();
            childGroups[child.name].Add(info);
        }
    }

    // Helper to generate a sanitized property name and ensure uniqueness.
    private static string SanitizePropertyName(string name, HashSet<string> usedNames)
    {
        string baseName = name.Replace(" ", "");
        if (string.IsNullOrEmpty(baseName))
            baseName = "child";
        string propName = char.ToLowerInvariant(baseName[0]) + baseName.Substring(1);
        string original = propName;
        int duplicateCount = 1;
        while (usedNames.Contains(propName))
        {
            propName = original + duplicateCount;
            duplicateCount++;
        }
        return propName;
    }

    private class ChildInfo
    {
        public string ChildName;
        public string TypeName;
        public ChildInfo(string childName, string typeName)
        {
            ChildName = childName;
            TypeName = typeName;
        }
    }

    // Returns a proper type string; if the type belongs to UnityEngine, return only the short name.
    private static string GetTypeString(Type type)
    {
        if (type.Namespace == "UnityEngine")
            return type.Name;
        return type.FullName;
    }

    // Sanitizes the class name by removing invalid characters and ensuring it starts with a letter.
    private static string SanitizeClassName(string name)
    {
        StringBuilder sb = new StringBuilder();
        foreach (char c in name)
            if (char.IsLetterOrDigit(c) || c == '_')
                sb.Append(c);
        string valid = sb.ToString();
        if (string.IsNullOrEmpty(valid) || !char.IsLetter(valid[0]))
            valid = "_" + valid;
        return valid;
    }
}

#endif