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;
        }
    }
}