ComponentConnector.cs

// Component Connector V2
// https://github.com/synchrok/ComponentConnector
// MIT License

using System;
#if UNITY_EDITOR
using System.Collections.Generic;
using System.Reflection;
using UnityEditor;
using UnityEditor.Events;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.Events;
using UnityEngine.SceneManagement;
using UnityEngine.UI;
#endif

public interface IComponentConnector { }

[AttributeUsage(AttributeTargets.Field)]
public class ComponentConnectAttribute : Attribute {
    public string name { get; }
    public bool inChildren { get; }

    public ComponentConnectAttribute(string name, bool inChildren = false) {
        this.name = name;
        this.inChildren = inChildren;
    }

    public ComponentConnectAttribute() {
        name = null;
        inChildren = true;
    }
}

[AttributeUsage(AttributeTargets.Method, AllowMultiple = true)]
public class OnClickAttribute : Attribute {
    public string name { get; }
    public bool inChildren { get; }

    public OnClickAttribute(string name, bool inChildren = false) {
        this.name = name;
        this.inChildren = inChildren;
    }
}

[AttributeUsage(AttributeTargets.Field)]
public class GetComponentAttribute : Attribute { }

/// <summary>
/// Legacy stub. 기존 씬/프리팹 호환용 — 새 시스템은 ComponentConnectorProcessor가 처리.
/// 씬/프리팹에서 이 컴포넌트를 제거한 뒤 이 클래스도 삭제 가능.
/// </summary>
public class ComponentConnector : UnityEngine.MonoBehaviour { }

#if UNITY_EDITOR
[InitializeOnLoad]
public static class ComponentConnectorProcessor {

    // ──────────────────────────────────────────────────────────────
    // Reflection Cache
    // ──────────────────────────────────────────────────────────────

    private struct ConnectFieldEntry {
        public FieldInfo Field;
        public string Name;
        public bool InChildren;
    }

    private struct OnClickEntry {
        public MethodInfo Method;
        public string ButtonName;
        public bool InChildren;
    }

    private struct GetComponentEntry {
        public FieldInfo Field;
    }

    private class CachedType {
        public ConnectFieldEntry[] ConnectFields;
        public OnClickEntry[] OnClickMethods;
        public GetComponentEntry[] GetComponentFields;
        public bool HasAny;
    }

    private static readonly Dictionary<Type, CachedType> _cache = new();
    private static bool _processScheduled;

    // ──────────────────────────────────────────────────────────────
    // Initialization
    // ──────────────────────────────────────────────────────────────

    static ComponentConnectorProcessor() {
        EditorApplication.hierarchyChanged += OnHierarchyChanged;
        PrefabStage.prefabStageOpened += OnPrefabStageOpened;
        PrefabStage.prefabSaving += OnPrefabSaving;
    }

    private static void OnHierarchyChanged() {
        if (Application.isPlaying) return;
        if (EditorApplication.isCompiling || EditorApplication.isUpdating) return;
        if (_processScheduled) return;

        _processScheduled = true;
        EditorApplication.delayCall += () => {
            _processScheduled = false;

            var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
            if (prefabStage != null)
                ProcessGameObject(prefabStage.prefabContentsRoot);
            else
                ProcessActiveScene();
        };
    }

    private static void OnPrefabStageOpened(PrefabStage stage) {
        ProcessGameObject(stage.prefabContentsRoot);
    }

    private static void OnPrefabSaving(GameObject root) {
        ProcessGameObject(root);
    }

    // ──────────────────────────────────────────────────────────────
    // Public API
    // ──────────────────────────────────────────────────────────────

    [MenuItem("CONTEXT/Component/Component Connect", false, int.MaxValue)]
    public static void ContextRun() {
        var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
        if (prefabStage != null)
            ProcessGameObject(prefabStage.prefabContentsRoot);
        else
            ProcessActiveScene();
    }

    public static void ProcessActiveScene() {
        var roots = SceneManager.GetActiveScene().GetRootGameObjects();
        foreach (var root in roots)
            ProcessGameObject(root);
    }

    public static void ProcessGameObject(GameObject root) {
        var components = root.GetComponentsInChildren<Component>(true);
        foreach (var comp in components) {
            if (comp == null) continue;

            var cached = GetOrBuild(comp.GetType());
            if (!cached.HasAny) continue;

            ProcessComponent(comp, cached);
        }
    }

    // ──────────────────────────────────────────────────────────────
    // Cache
    // ──────────────────────────────────────────────────────────────

    private static CachedType GetOrBuild(Type type) {
        if (_cache.TryGetValue(type, out var cached))
            return cached;

        cached = Build(type);
        _cache[type] = cached;
        return cached;
    }

    private static CachedType Build(Type type) {
        var ct = new CachedType();
        var isCC = typeof(IComponentConnector).IsAssignableFrom(type);

        var connects = new List<ConnectFieldEntry>();
        var getComps = new List<GetComponentEntry>();
        var clicks = new List<OnClickEntry>();

        var fields = type.GetFields(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance);
        foreach (var fi in fields) {
            foreach (var attr in fi.GetCustomAttributes(true)) {
                if (attr is ComponentConnectAttribute cca)
                    connects.Add(new ConnectFieldEntry {
                        Field = fi,
                        Name = cca.name ?? fi.Name,
                        InChildren = cca.inChildren
                    });
                else if (attr is GetComponentAttribute)
                    getComps.Add(new GetComponentEntry { Field = fi });
            }
        }

        if (isCC) {
            var methods = type.GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance);
            foreach (var mi in methods) {
                foreach (var attr in mi.GetCustomAttributes(typeof(OnClickAttribute), true)) {
                    if (attr is OnClickAttribute oca)
                        clicks.Add(new OnClickEntry {
                            Method = mi,
                            ButtonName = oca.name,
                            InChildren = oca.inChildren
                        });
                }
            }
        }

        ct.ConnectFields = connects.ToArray();
        ct.GetComponentFields = getComps.ToArray();
        ct.OnClickMethods = clicks.ToArray();
        ct.HasAny = connects.Count > 0 || getComps.Count > 0 || clicks.Count > 0;

        return ct;
    }

    // ──────────────────────────────────────────────────────────────
    // Processing
    // ──────────────────────────────────────────────────────────────

    private static void ProcessComponent(Component mb, CachedType cached) {
        bool changed = false;
        bool undoRecorded = false;

        foreach (var entry in cached.ConnectFields) {
            if (ShouldSkip(entry.Field, mb)) continue;
            if (TryConnect(mb, entry, ref undoRecorded))
                changed = true;
        }

        foreach (var entry in cached.GetComponentFields) {
            if (ShouldSkip(entry.Field, mb)) continue;
            if (TryGetComp(mb, entry, ref undoRecorded))
                changed = true;
        }

        foreach (var entry in cached.OnClickMethods) {
            if (TryOnClick(mb, entry))
                changed = true;
        }

        if (changed)
            EditorUtility.SetDirty(mb);
    }

    // ──────────────────────────────────────────────────────────────
    // Skip Logic
    // ──────────────────────────────────────────────────────────────

    private static bool ShouldSkip(FieldInfo fi, Component mb) {
        var value = fi.GetValue(mb);
        if (value == null) return false;

        if (value is UnityEngine.Object obj)
            return obj != null; // alive → skip; destroyed → reconnect

        if (fi.FieldType.IsArray)
            return value is Array arr && arr.Length > 0;

        return true;
    }

    // ──────────────────────────────────────────────────────────────
    // [ComponentConnect]
    // ──────────────────────────────────────────────────────────────

    private static bool TryConnect(Component mb, ConnectFieldEntry entry, ref bool undoRecorded) {
        var name = entry.Name;
        var inChildren = entry.InChildren;

        if (name.EndsWith("~"))
            return TryConnectArray(mb, entry.Field, name.TrimEnd('~'), inChildren, ref undoRecorded);

        GameObject go;

        if (name.Contains("..."))
            go = ResolveEllipsis(mb, name, inChildren);
        else
            go = ResolveStandard(mb, name, inChildren);

        if (go == null) return false;

        object newValue;
        if (entry.Field.FieldType == typeof(GameObject))
            newValue = go;
        else {
            var comp = go.GetComponent(entry.Field.FieldType);
            if (comp == null) return false;
            newValue = comp;
        }

        RecordUndo(mb, ref undoRecorded);
        entry.Field.SetValue(mb, newValue);
        return true;
    }

    private static bool TryConnectArray(Component mb, FieldInfo fi, string keyword,
                                         bool inChildren, ref bool undoRecorded) {
        var searched = CCSearchAllStartsWith(mb.gameObject, keyword);
        if (searched == null && !inChildren)
            searched = CCSearchAllStartsWith(mb.transform.root.gameObject, keyword);
        if (searched == null)
            searched = CCSearchAllStartsWithSceneRoots(keyword);

        if (searched == null) return false;

        RecordUndo(mb, ref undoRecorded);

        if (fi.FieldType == typeof(GameObject[])) {
            fi.SetValue(mb, searched);
        } else {
            var elemType = fi.FieldType.GetElementType();
            var arr = Array.CreateInstance(elemType, searched.Length);
            for (var i = 0; i < searched.Length; i++)
                arr.SetValue(searched[i].GetComponent(elemType), i);
            fi.SetValue(mb, arr);
        }

        return true;
    }

    private static GameObject ResolveStandard(Component mb, string name, bool inChildren) {
        var go = CCSearch(mb.gameObject, name);
        if (go == null && !inChildren)
            go = CCSearch(mb.transform.root.gameObject, name);
        if (go == null)
            go = CCSearchSceneRoots(name, mb.transform.root.gameObject);
        return go;
    }

    private static GameObject ResolveEllipsis(Component mb, string ccName, bool inChildren) {
        var d = ccName.Split(new[] { "..." }, StringSplitOptions.RemoveEmptyEntries);

        if (d.Length == 1) {
            var go = CCSearch(mb.gameObject, d[0]);
            if (go == null && !inChildren)
                go = CCSearch(mb.transform.root.gameObject, d[0]);
            if (go == null)
                go = CCSearchSceneRoots(d[0], mb.transform.root.gameObject);
            return go;
        }

        var result = CCSearchWithParentName(mb.gameObject, d[1], d[0]);
        if (result == null && !inChildren)
            result = CCSearchWithParentName(mb.transform.root.gameObject, d[1], d[0]);
        if (result == null) {
            var roots = SceneManager.GetActiveScene().GetRootGameObjects();
            foreach (var root in roots) {
                if (root == mb.transform.root.gameObject) continue;
                result = CCSearchWithParentName(root, d[1], d[0]);
                if (result != null) break;
            }
        }

        return result;
    }

    // ──────────────────────────────────────────────────────────────
    // [GetComponent]
    // ──────────────────────────────────────────────────────────────

    private static bool TryGetComp(Component mb, GetComponentEntry entry, ref bool undoRecorded) {
        if (mb.gameObject == null) return false;

        object newValue;
        if (entry.Field.FieldType == typeof(GameObject))
            newValue = mb.gameObject;
        else {
            var comp = mb.gameObject.GetComponent(entry.Field.FieldType);
            if (comp == null) return false;
            newValue = comp;
        }

        RecordUndo(mb, ref undoRecorded);
        entry.Field.SetValue(mb, newValue);
        return true;
    }

    // ──────────────────────────────────────────────────────────────
    // [OnClick]
    // ──────────────────────────────────────────────────────────────

    private static bool TryOnClick(Component mb, OnClickEntry entry) {
        var go = CCSearch(mb.gameObject, entry.ButtonName);
        if (go == null && !entry.InChildren)
            go = CCSearch(mb.transform.root.gameObject, entry.ButtonName);

        if (go == null) return false;

        var btn = go.GetComponent<Button>();
        if (btn == null) return false;

        var pcnt = btn.onClick.GetPersistentEventCount();
        if (pcnt == 1 &&
            btn.onClick.GetPersistentTarget(0) == (UnityEngine.Object)mb &&
            btn.onClick.GetPersistentMethodName(0) == entry.Method.Name)
            return false;

        Undo.RecordObject(btn, "ComponentConnector");

        for (var i = pcnt - 1; i >= 0; i--)
            UnityEventTools.RemovePersistentListener(btn.onClick, 0);

        var action = Delegate.CreateDelegate(typeof(UnityAction), mb, entry.Method.Name) as UnityAction;
        UnityEventTools.AddPersistentListener(btn.onClick, action);

        EditorUtility.SetDirty(btn);
        return true;
    }

    // ──────────────────────────────────────────────────────────────
    // Undo
    // ──────────────────────────────────────────────────────────────

    private static void RecordUndo(Component mb, ref bool undoRecorded) {
        if (undoRecorded) return;
        Undo.RecordObject(mb, "ComponentConnector");
        undoRecorded = true;
    }

    // ──────────────────────────────────────────────────────────────
    // CC Search (Private)
    // ──────────────────────────────────────────────────────────────

    private static GameObject CCSearch(GameObject target, string name) {
        if (string.Equals(target.name, name, StringComparison.OrdinalIgnoreCase))
            return target;

        for (var i = 0; i < target.transform.childCount; i++) {
            var result = CCSearch(target.transform.GetChild(i).gameObject, name);
            if (result != null) return result;
        }

        return null;
    }

    private static GameObject CCSearchWithParentName(GameObject target, string name, string parentName) {
        if (string.Equals(target.name, parentName, StringComparison.OrdinalIgnoreCase)) {
            for (var i = 0; i < target.transform.childCount; i++) {
                var child = target.transform.GetChild(i);
                if (child.name.Equals(name))
                    return child.gameObject;
            }
        }

        for (var i = 0; i < target.transform.childCount; i++) {
            var result = CCSearchWithParentName(target.transform.GetChild(i).gameObject, name, parentName);
            if (result != null) return result;
        }

        return null;
    }

    private static readonly List<GameObject> _ccSearchList = new();

    private static GameObject[] CCSearchAllStartsWith(GameObject target, string keyword) {
        _ccSearchList.Clear();
        CCSearchAllStartsWithImpl(target, keyword);
        return _ccSearchList.Count > 0 ? _ccSearchList.ToArray() : null;
    }

    private static void CCSearchAllStartsWithImpl(GameObject target, string keyword) {
        if (target.name.StartsWith(keyword)) {
            _ccSearchList.Add(target);
            return;
        }

        for (var i = 0; i < target.transform.childCount; i++)
            CCSearchAllStartsWithImpl(target.transform.GetChild(i).gameObject, keyword);
    }

    private static GameObject CCSearchSceneRoots(string name, GameObject exclude) {
        var roots = SceneManager.GetActiveScene().GetRootGameObjects();
        foreach (var root in roots) {
            if (root == exclude) continue;
            var go = CCSearch(root, name);
            if (go != null) return go;
        }
        return null;
    }

    private static GameObject[] CCSearchAllStartsWithSceneRoots(string keyword) {
        var roots = SceneManager.GetActiveScene().GetRootGameObjects();
        foreach (var root in roots) {
            var result = CCSearchAllStartsWith(root, keyword);
            if (result != null) return result;
        }
        return null;
    }
}
#endif