TheInitialProject/Assets/CaoCao/Scripts/Editor/XAsset/GUI/TreeViews/BuildAssetTreeView.cs

using CaoCao.XAsset;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using UnityEditor;
using UnityEditor.IMGUI.Controls;
using UnityEngine;
using Object = UnityEngine.Object;

namespace CaoCao.Editor
{
    internal sealed class BuildAssetTreeViewItem : TreeViewItem
    {
        public readonly BuildAsset data;

        public BuildAssetTreeViewItem(BuildAsset asset, int depth) : base(asset.path.GetHashCode(), depth, asset.path)
        {
            icon = AssetDatabase.GetCachedIcon(asset.path) as Texture2D;
            data = asset;
        }
    }

    internal class BuildAssetTreeView : TreeView
    {
        public readonly List<BuildAsset> assets = new List<BuildAsset>();

        private readonly SortOption[] m_SortOptions =
            {SortOption.Asset, SortOption.Size, SortOption.Type, SortOption.Bundle, SortOption.AddressMode};

        private readonly List<TreeViewItem> result = new List<TreeViewItem>();

        public BuildAssetTreeView(TreeViewState treeViewState, MultiColumnHeaderState headerState) :
            base(treeViewState,
                new MultiColumnHeader(headerState))
        {
            showBorder = true;
            showAlternatingRowBackgrounds = true;
            multiColumnHeader.sortingChanged += OnSortingChanged;
            multiColumnHeader.ResizeToFit();
        }

        private void OnSortingChanged(MultiColumnHeader header)
        {
            SortIfNeeded(rootItem, GetRows());
        }

        private void SortIfNeeded(TreeViewItem root, IList<TreeViewItem> rows)
        {
            if (rows.Count <= 1) return;

            if (multiColumnHeader.sortedColumnIndex == -1) return;

            SortByColumn();

            rows.Clear();

            foreach (var t in root.children)
            {
                rows.Add(t);
                if (!t.hasChildren || t.children[0] == null || !IsExpanded(t.id)) continue;

                foreach (var child in t.children) rows.Add(child);
            }

            Repaint();
        }

        internal static MultiColumnHeaderState CreateDefaultMultiColumnHeaderState()
        {
            return new MultiColumnHeaderState(GetColumns());
        }

        private static MultiColumnHeaderState.Column[] GetColumns()
        {
            var retVal = new[]
            {
                new MultiColumnHeaderState.Column
                {
                    headerContent = new GUIContent("Asset"),
                    minWidth = 64,
                    width = 200,
                    headerTextAlignment = TextAlignment.Left,
                    canSort = true,
                    autoResize = true
                },
                new MultiColumnHeaderState.Column
                {
                    headerContent = new GUIContent("Bundle"),
                    minWidth = 64,
                    width = 96,
                    headerTextAlignment = TextAlignment.Left,
                    canSort = true,
                    autoResize = true
                },
                new MultiColumnHeaderState.Column
                {
                    headerContent = new GUIContent("Size"),
                    minWidth = 64,
                    width = 96,
                    headerTextAlignment = TextAlignment.Left,
                    canSort = true,
                    autoResize = true
                },
                new MultiColumnHeaderState.Column
                {
                    headerContent = new GUIContent("Type"),
                    minWidth = 64,
                    width = 96,
                    headerTextAlignment = TextAlignment.Left,
                    canSort = true,
                    autoResize = true
                },
                new MultiColumnHeaderState.Column
                {
                    headerContent = new GUIContent("Address Mode"),
                    minWidth = 64,
                    width = 96,
                    headerTextAlignment = TextAlignment.Left,
                    canSort = true,
                    autoResize = true
                }
            };
            return retVal;
        }

        protected override void ContextClickedItem(int id)
        {
            base.ContextClickedItem(id);
            var items = Array.ConvertAll(GetSelection().ToArray(), o => (BuildAssetTreeViewItem) FindItem(o, rootItem));
            var menu = new GenericMenu();
            menu.AddItem(new GUIContent("Copy Asset Path"), false,
                data =>
                {
                    EditorGUIUtility.systemCopyBuffer =
                        string.Join("\n", Array.ConvertAll(items, input => input.displayName));
                }, null);

            menu.AddItem(new GUIContent("Remove Bundles"), false, data =>
            {
                var assetPaths = Array.ConvertAll(items, input => input.displayName);
                Builder.RemoveBundles(assetPaths);
            }, null);

            var configs = Settings.FindAssets<AssetPack>();
            foreach (var config in configs)
                menu.AddItem(new GUIContent($"Export to...AssetPack/{config.name}"), false,
                    data =>
                    {
                        var set = new HashSet<Object>();
                        if (config.assets != null) set.UnionWith(config.assets);
                        foreach (var item in items) set.Add(AssetDatabase.LoadAssetAtPath<Object>(item.data.path));
                        config.assets = set.ToArray();
                        Selection.activeObject = config;
                        EditorUtility.SetDirty(config);
                        EditorUtility.FocusProjectWindow();
                        AssetDatabase.SaveAssets();
                    }, null);
            menu.ShowAsContext();
        }


        protected override TreeViewItem BuildRoot()
        {
            var root = new TreeViewItem {id = 0, depth = -1, displayName = "Root"};
            foreach (var asset in assets)
            {
                var childItem = new BuildAssetTreeViewItem(asset, root.depth + 1);
                root.AddChild(childItem);
            }

            return root;
        }

        public override void OnGUI(Rect rect)
        {
            base.OnGUI(rect);
            if (UnityEngine.Event.current.type == EventType.MouseDown && UnityEngine.Event.current.button == 0 &&
                rect.Contains(UnityEngine.Event.current.mousePosition))
                SetSelection(Array.Empty<int>(), TreeViewSelectionOptions.FireSelectionChanged);
        }

        protected override IList<TreeViewItem> BuildRows(TreeViewItem root)
        {
            var rows = (List<TreeViewItem>) base.BuildRows(root);
            if (!string.IsNullOrEmpty(searchString))
            {
                result.Clear();
                var stack = new Stack<TreeViewItem>();
                foreach (var element in root.children) stack.Push(element);

                while (stack.Count > 0)
                {
                    var current = stack.Pop();
                    // Matches search?
                    if (current.displayName.IndexOf(searchString, StringComparison.OrdinalIgnoreCase) >= 0)
                        result.Add(current);

                    // if (current.children != null && current.children.Count > 0)
                    // {
                    //     foreach (var element in current.children)
                    //     {
                    //         stack.Push(element);
                    //     }
                    // }
                }

                rows = result;
            }

            SortIfNeeded(root, rows);
            return rows;
        }

        private void SortIfNeeded(TreeViewItem root, List<TreeViewItem> rows)
        {
            if (rows.Count <= 1) return;

            if (multiColumnHeader.sortedColumnIndex == -1) return;

            SortByColumn();

            rows.Clear();

            foreach (var t in root.children)
            {
                rows.Add(t);
                if (!t.hasChildren || t.children[0] == null || !IsExpanded(t.id)) continue;

                foreach (var child in t.children) rows.Add(child);
            }

            Repaint();
        }

        private void SortByColumn()
        {
            var sortedColumns = multiColumnHeader.state.sortedColumns;

            if (sortedColumns.Length == 0) return;

            var assetList = new List<TreeViewItem>();
            foreach (var item in rootItem.children) assetList.Add(item);

            var orderedItems = InitialOrder(assetList, sortedColumns);
            rootItem.children = orderedItems.ToList();
        }

        private IEnumerable<TreeViewItem> InitialOrder(IEnumerable<TreeViewItem> myTypes, int[] columnList)
        {
            var sortOption = m_SortOptions[columnList[0]];
            var ascending = multiColumnHeader.IsSortedAscending(columnList[0]);
            switch (sortOption)
            {
                case SortOption.Asset:
                    return myTypes.Order(l => l.displayName, ascending);
                case SortOption.Bundle:
                    return myTypes.Order(l => ((BuildAssetTreeViewItem) l).data.bundle, ascending);
                case SortOption.Size:
                    return myTypes.Order(l => ((BuildAssetTreeViewItem) l).data.size, ascending);
                case SortOption.Type:
                    return myTypes.Order(l => ((BuildAssetTreeViewItem) l).data.type, ascending);
                case SortOption.AddressMode:
                    return myTypes.Order(l => ((BuildAssetTreeViewItem) l).data.addressMode, ascending);
                default:
                    return myTypes.Order(l => l.displayName, ascending);
            }
        }

        protected override void RowGUI(RowGUIArgs args)
        {
            for (var i = 0; i < args.GetNumVisibleColumns(); ++i)
            {
                if (args.item is BuildAssetTreeViewItem item)
                {
                    if (item.data == null)
                        using (new EditorGUI.DisabledScope(true))
                        {
                            base.RowGUI(args);
                        }
                    else
                        CellGUI(args.GetCellRect(i), item, args.GetColumn(i), ref args);
                }
                else
                {
                    base.RowGUI(args);
                }
            }
        }

        private void CellGUI(Rect cellRect, BuildAssetTreeViewItem item, int column, ref RowGUIArgs args)
        {
            CenterRectUsingSingleLineHeight(ref cellRect);

            switch (column)
            {
                case 0:
                    cellRect.xMin += GetContentIndent(item) + extraSpaceBeforeIconAndLabel;
                    var iconRect = new Rect(cellRect.x + 1, cellRect.y + 1, cellRect.height - 2, cellRect.height - 2);
                    if (item.icon != null) GUI.DrawTexture(iconRect, item.icon, ScaleMode.ScaleToFit);

                    DefaultGUI.Label(
                        new Rect(cellRect.x + iconRect.xMax + 1, cellRect.y, cellRect.width - iconRect.width,
                            cellRect.height),
                        item.displayName,
                        args.selected,
                        args.focused);
                    break;
                case 1:
                    if (item.data.addressMode != AddressMode.LoadByDependencies) item.data.bundle = Settings.PackAsset(item.data);
                    DefaultGUI.Label(cellRect, item.data.bundle, args.selected, args.focused);
                    break;
                case 2:
                    DefaultGUI.Label(cellRect, CaoCao.XAsset.Utility.FormatBytes(item.data.size), args.selected,
                        args.focused);
                    break;
                case 3:
                    DefaultGUI.Label(cellRect, item.data.type, args.selected,
                        args.focused);
                    break;
                case 4:
                    DefaultGUI.Label(cellRect, item.data.addressMode.ToString(), args.selected,
                        args.focused);
                    break;
            }
        }

        protected override void DoubleClickedItem(int id)
        {
            var assetItem = FindItem(id, rootItem);
            if (assetItem == null) return;

            var o = AssetDatabase.LoadAssetAtPath<Object>(assetItem.displayName);
            EditorGUIUtility.PingObject(o);
            Selection.activeObject = o;
        }

        protected override bool CanBeParent(TreeViewItem item)
        {
            return true;
        }

        protected override bool CanStartDrag(CanStartDragArgs args)
        {
            args.draggedItemIDs = GetSelection();
            return DragAndDrop.paths.Length != 0;
        }

        protected override void SetupDragAndDrop(SetupDragAndDropArgs args)
        {
        }

        protected override DragAndDropVisualMode HandleDragAndDrop(DragAndDropArgs args)
        {
            return DragAndDropVisualMode.Rejected;
        }

        private enum SortOption
        {
            Asset,
            Size,
            Type,
            Bundle,
            AddressMode
        }
    }
}