Editor/AssetBrowserAddon/LocationRegistration.cs
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using Editor;
using Sandbox;
namespace Sandbox.AssetBrowserAddon;
internal static partial class LocationRegistration
{
private const string CookieKey = "Spud.CustomLocations";
private static readonly List<CustomLocationDefinition> _definitions =
ProjectCookie.Get(CookieKey, new List<CustomLocationDefinition>()) ?? new();
private static Type _folderNodeType;
private static bool _sceneAttached;
private static bool _hammerAttached;
private sealed class TreeContext
{
public AssetLocations Tree;
public IList Items;
public CustomLocationsHeader Header;
public readonly TreeNode.Spacer Spacer = new(10);
}
private static readonly List<TreeContext> _trees = new();
[EditorEvent.Frame]
private static void TryInstall()
{
_folderNodeType ??= typeof(LocalAssetBrowser).Assembly.GetType("Editor.FolderNode");
if (_folderNodeType is null)
return;
if (_sceneAttached)
return;
if (MainAssetBrowser.Instance?.Local is not LocalAssetBrowser local)
return;
if (RegisterBrowser(local))
{
_sceneAttached = true;
RebuildAll();
}
}
internal static bool RegisterBrowser(LocalAssetBrowser browser)
{
if (browser is null) return false;
var tree = GetAssetLocationsShared(browser);
if (tree is null) return false;
if (_trees.Any(x => ReferenceEquals(x.Tree, tree))) return false;
var ctx = new TreeContext
{
Tree = tree,
Items = ResolveItemList(tree),
Header = new CustomLocationsHeader { OnAddRequested = () => ShowLocationDialog() }
};
AttachHeader(ctx);
_trees.Add(ctx);
return true;
}
private static void AttachHeader(TreeContext ctx)
{
if (ctx.Tree is null)
return;
if (ctx.Items is null)
{
if (!ContainsNode(ctx.Tree.Items, ctx.Header))
{
ctx.Tree.AddItem(ctx.Spacer);
ctx.Tree.AddItem(ctx.Header);
}
}
else
{
if (!ContainsNode(ctx.Items, ctx.Header))
{
var insertIndex = Math.Min(2, ctx.Items.Count);
if (!ContainsNode(ctx.Items, ctx.Spacer))
{
ctx.Items.Insert(insertIndex, ctx.Spacer);
insertIndex++;
}
ctx.Items.Insert(insertIndex, ctx.Header);
}
}
}
private static IList ResolveItemList(AssetLocations tree)
{
var type = tree.GetType();
while (type is not null)
{
var field = type.GetField("_items", BindingFlags.Instance | BindingFlags.NonPublic);
if (field?.GetValue(tree) is IList list)
return list;
type = type.BaseType;
}
return null;
}
private static void RebuildCustomNodes(TreeContext ctx)
{
if (ctx?.Header is null || _folderNodeType is null)
return;
ctx.Header.Clear();
foreach (var location in BuildBuiltInLocations())
{
var node = CreateNode(location, null);
if (node is not null)
ctx.Header.AddItem(node);
}
foreach (var definition in _definitions)
{
var node = CreateNode(new FilteredLocation(definition), definition);
if (node is not null)
ctx.Header.AddItem(node);
}
ctx.Header.Dirty();
ctx.Tree?.RefreshChildren();
ctx.Tree?.Open(ctx.Header);
}
private static TreeNode CreateNode(LocalAssetBrowser.Location location, CustomLocationDefinition definition)
{
if (Activator.CreateInstance(_folderNodeType, location) is not TreeNode node)
return null;
if (definition is not null)
{
var menuField = _folderNodeType.GetField("OnContextMenuOpen",
BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic);
menuField?.SetValue(node, (Action)(() => ShowCustomLocationMenu(definition)));
}
return node;
}
private static void ShowCustomLocationMenu(CustomLocationDefinition definition)
{
var menu = new ContextMenu();
menu.AddOption("Edit...", "edit", () => ShowLocationDialog(definition));
menu.AddOption("Remove", "delete", () => RemoveDefinition(definition));
menu.OpenAt(Editor.Application.CursorPosition);
}
private static void ShowLocationDialog(CustomLocationDefinition definition = null)
{
var dialog = new LocationEditor(updated => OnLocationSaved(definition, updated), definition);
dialog.Show();
}
private static void OnLocationSaved(CustomLocationDefinition original, CustomLocationDefinition definition)
{
if (definition is null)
return;
if (original is null)
{
_definitions.Add(definition);
}
else
{
var index = _definitions.FindIndex(x => x.Id == original.Id);
if (index >= 0)
_definitions[index] = definition;
else
_definitions.Add(definition);
}
SaveDefinitions();
RebuildAll();
}
private static void RemoveDefinition(CustomLocationDefinition definition)
{
if (definition is null)
return;
_definitions.RemoveAll(x => x.Id == definition.Id);
SaveDefinitions();
RebuildAll();
}
private static IEnumerable<LocalAssetBrowser.Location> BuildBuiltInLocations()
{
yield return new MaterialLocation();
yield return new ModelLocation();
}
private static void SaveDefinitions()
{
ProjectCookie.Set(CookieKey, _definitions);
}
internal static AssetLocations GetAssetLocationsShared(LocalAssetBrowser browser)
{
var assetLocationsField = typeof(AssetBrowser).GetField("AssetLocations",
BindingFlags.Instance | BindingFlags.NonPublic);
return assetLocationsField?.GetValue(browser) as AssetLocations;
}
internal static void RebuildAll()
{
foreach (var ctx in _trees)
{
RebuildCustomNodes(ctx);
}
}
private static bool ContainsNode(IEnumerable collection, object value)
{
foreach (var item in collection)
{
if (ReferenceEquals(item, value))
return true;
}
return false;
}
private sealed class CustomLocationsHeader : TreeNode.SmallHeader
{
public Action OnAddRequested { get; set; }
public CustomLocationsHeader() : base("bookmarks", "Bookmarks")
{
}
public override bool OnContextMenu()
{
var menu = new ContextMenu();
menu.AddOption("Add Bookmark...", "add", () => OnAddRequested?.Invoke());
menu.OpenAt( Editor.Application.CursorPosition );
return true;
}
}
}