@* Reusable hierarchical tree view with expand/collapse, ARIA roles, and guide lines *@
@typeparam TItem
@inject IJSRuntime JSRuntime

@if (_items is null || _items.Count == 0)
{
    if (EmptyContent != null)
    {
        @EmptyContent
    }
}
else
{
    <ul role="tree" class="tv-root @(ShowGuideLines ? "tv-guides" : "")">
        @foreach (var item in _items)
        {
            RenderNode(item, 0);
        }
    </ul>
}

@if (_showContextMenu && _contextMenuItem != null && ContextMenu != null)
{
    <div class="tv-ctx-overlay" @onclick="DismissContextMenu" style="position:fixed;top:0;left:0;width:100vw;height:100vh;z-index:1049;background:transparent;"></div>
    <div class="dropdown-menu show" tabindex="-1" @ref="_contextMenuRef" @onkeydown="OnContextMenuKeyDown"
         @onclick="DismissContextMenu"
         style="position:fixed;top:@(_contextMenuY)px;left:@(_contextMenuX)px;z-index:1050;outline:none;">
        @ContextMenu(_contextMenuItem)
    </div>
}

@{ void RenderNode(TItem item, int depth)
    {
        var key = KeySelector(item);
        var children = ChildrenSelector(item);
        var isBranch = HasChildrenSelector(item);
        var isExpanded = _expandedKeys.Contains(KeyStr(key));
        var isSelected = Selectable && SelectedKey != null && SelectedKey.Equals(key);
        var rowClasses = "tv-row" + (isSelected ? " tv-selected " + SelectedCssClass : "");

        <li role="treeitem" @key="key"
            aria-expanded="@(isBranch ? (isExpanded ? "true" : "false") : null)"
            aria-selected="@(isSelected ? "true" : null)">
            <div class="@rowClasses" style="padding-left: @(depth * IndentPx)px; --tv-depth: @depth;"
                @oncontextmenu="(e) => OnContextMenu(e, item)" @oncontextmenu:preventDefault="@(ContextMenu != null)">
                @if (isBranch)
                {
                    <span class="tv-toggle" @onclick="() => ToggleExpand(key)" @onclick:stopPropagation><i class="bi bi-chevron-right"></i></span>
                }
                else
                {
                    <span class="tv-spacer"></span>
                }
                <span class="tv-content" @onclick="() => OnContentClick(key)" @onclick:stopPropagation>
                    @NodeContent(item)
                </span>
            </div>
            @if (isBranch && isExpanded && children is { Count: > 0 })
            {
                <ul role="group">
                    @foreach (var child in children)
                    {
                        RenderNode(child, depth + 1);
                    }
                </ul>
            }
        </li>
    }
}

@code {
    private IReadOnlyList<TItem>? _items;
    private HashSet<string> _expandedKeys = new();

    /// <summary>Normalize any key object to a string for consistent comparison with sessionStorage.</summary>
    private string KeyStr(object key) => key.ToString()!;
    private bool _initialExpansionApplied;
    private bool _storageLoaded;
    private TItem? _contextMenuItem;
    private double _contextMenuX;
    private double _contextMenuY;
    private bool _showContextMenu;
    private bool _contextMenuNeedsFocus;
    private ElementReference _contextMenuRef;

    [Parameter, EditorRequired] public IReadOnlyList<TItem> Items { get; set; } = [];
    [Parameter, EditorRequired] public Func<TItem, IReadOnlyList<TItem>> ChildrenSelector { get; set; } = default!;
    [Parameter, EditorRequired] public Func<TItem, bool> HasChildrenSelector { get; set; } = default!;
    [Parameter, EditorRequired] public Func<TItem, object> KeySelector { get; set; } = default!;
    [Parameter, EditorRequired] public RenderFragment<TItem> NodeContent { get; set; } = default!;
    [Parameter] public RenderFragment? EmptyContent { get; set; }
    [Parameter] public RenderFragment<TItem>? ContextMenu { get; set; }
    [Parameter] public int IndentPx { get; set; } = 24;
    [Parameter] public bool ShowGuideLines { get; set; } = true;
    [Parameter] public Func<TItem, bool>? InitiallyExpanded { get; set; }
    [Parameter] public bool Selectable { get; set; }
    [Parameter] public object? SelectedKey { get; set; }
    [Parameter] public EventCallback<object?> SelectedKeyChanged { get; set; }
    [Parameter] public string SelectedCssClass { get; set; } = "bg-primary bg-opacity-10";
    [Parameter] public string? StorageKey { get; set; }

    protected override void OnParametersSet()
    {
        _items = Items;

        if (!_initialExpansionApplied && InitiallyExpanded != null && _items is { Count: > 0 })
        {
            // Only apply InitiallyExpanded when there is no StorageKey, or storage
            // has already been checked and returned nothing (no prior state).
            if (StorageKey == null || (_storageLoaded && _expandedKeys.Count == 0))
            {
                _initialExpansionApplied = true;
                ApplyInitialExpansion(_items);
            }
        }

        // Clear selection if the selected key no longer exists in the current items tree
        if (Selectable && SelectedKey != null && _items is not null && !KeyExistsInTree(_items, SelectedKey))
        {
            _ = SelectedKeyChanged.InvokeAsync(null);
        }
    }

    protected override async Task OnAfterRenderAsync(bool firstRender)
    {
        if (_contextMenuNeedsFocus && _showContextMenu)
        {
            _contextMenuNeedsFocus = false;
            try { await _contextMenuRef.FocusAsync(); } catch { /* element may have been disposed if dismissed during render */ }
        }

        if (firstRender && StorageKey != null)
        {
            var json = await JSRuntime.InvokeAsync<string?>("treeviewStorage.load", StorageKey);
            _storageLoaded = true;

            if (json != null)
            {
                var keys = System.Text.Json.JsonSerializer.Deserialize<List<string>>(json);
                if (keys != null)
                {
                    // Union (don't replace): callers may have invoked RevealNode before
                    // this async storage load completed. Preserving those reveal-added
                    // keys ensures deep-link reveal isn't clobbered by the restore.
                    foreach (var k in keys) _expandedKeys.Add(k);
                    _initialExpansionApplied = true;
                }
            }
            else if (InitiallyExpanded != null && _items is { Count: > 0 } && !_initialExpansionApplied)
            {
                // Storage returned null (no prior state) — fall back to InitiallyExpanded
                _initialExpansionApplied = true;
                ApplyInitialExpansion(_items);
            }

            StateHasChanged();
        }
    }

    private bool KeyExistsInTree(IReadOnlyList<TItem> items, object key)
    {
        foreach (var item in items)
        {
            if (key.Equals(KeySelector(item)))
                return true;

            var children = ChildrenSelector(item);
            if (children is { Count: > 0 } && KeyExistsInTree(children, key))
                return true;
        }
        return false;
    }

    private void ApplyInitialExpansion(IReadOnlyList<TItem> items)
    {
        foreach (var item in items)
        {
            if (InitiallyExpanded!(item))
            {
                _expandedKeys.Add(KeyStr(KeySelector(item)));
            }

            var children = ChildrenSelector(item);
            if (children is { Count: > 0 })
            {
                ApplyInitialExpansion(children);
            }
        }
    }

    private void ToggleExpand(object key)
    {
        var k = KeyStr(key);
        if (!_expandedKeys.Remove(k))
        {
            _expandedKeys.Add(k);
        }

        PersistExpandedState();
    }

    private void PersistExpandedState()
    {
        if (StorageKey != null)
        {
            var json = System.Text.Json.JsonSerializer.Serialize(_expandedKeys.ToList());
            _ = JSRuntime.InvokeVoidAsync("treeviewStorage.save", StorageKey, json);
        }
    }

    private async Task OnContentClick(object key)
    {
        if (Selectable)
        {
            await SelectedKeyChanged.InvokeAsync(key);
        }
    }

    private void OnContextMenu(MouseEventArgs e, TItem item)
    {
        if (ContextMenu == null) return;

        _contextMenuItem = item;
        _contextMenuX = e.ClientX;
        _contextMenuY = e.ClientY;
        _showContextMenu = true;
        _contextMenuNeedsFocus = true;
    }

    private void DismissContextMenu()
    {
        _showContextMenu = false;
        _contextMenuItem = default;
    }

    private void OnContextMenuKeyDown(KeyboardEventArgs e)
    {
        if (e.Key == "Escape")
        {
            DismissContextMenu();
        }
    }

    /// <summary>Whether the node with the given key is currently expanded.</summary>
    public bool IsExpanded(object key) => _expandedKeys.Contains(KeyStr(key));

    /// <summary>Expand every branch node in the tree.</summary>
    public void ExpandAll()
    {
        if (_items is { Count: > 0 })
        {
            ExpandAllRecursive(_items);
        }

        PersistExpandedState();
        StateHasChanged();
    }

    private void ExpandAllRecursive(IReadOnlyList<TItem> items)
    {
        foreach (var item in items)
        {
            if (HasChildrenSelector(item))
            {
                _expandedKeys.Add(KeyStr(KeySelector(item)));
            }

            var children = ChildrenSelector(item);
            if (children is { Count: > 0 })
            {
                ExpandAllRecursive(children);
            }
        }
    }

    /// <summary>Collapse every node in the tree.</summary>
    public void CollapseAll()
    {
        _expandedKeys.Clear();
        PersistExpandedState();
        StateHasChanged();
    }

    /// <summary>
    /// Expand all ancestors of the given key so it becomes visible.
    /// Optionally select the node.
    /// </summary>
    public async Task RevealNode(object key, bool select = false)
    {
        var parentLookup = BuildParentLookup();
        var k = KeyStr(key);

        // If key is not in the tree at all, no-op
        if (!parentLookup.ContainsKey(k))
            return;

        // Walk up through ancestors
        var current = k;
        while (parentLookup.TryGetValue(current, out var parentKey) && parentKey != null)
        {
            _expandedKeys.Add(parentKey);
            current = parentKey;
        }

        if (select && Selectable)
        {
            await SelectedKeyChanged.InvokeAsync(key);
        }

        PersistExpandedState();
        StateHasChanged();
    }

    private Dictionary<string, string?> BuildParentLookup()
    {
        var lookup = new Dictionary<string, string?>();
        if (_items is { Count: > 0 })
        {
            BuildParentLookupRecursive(_items, null, lookup);
        }
        return lookup;
    }

    private void BuildParentLookupRecursive(IReadOnlyList<TItem> items, string? parentKey, Dictionary<string, string?> lookup)
    {
        foreach (var item in items)
        {
            var key = KeyStr(KeySelector(item));
            lookup[key] = parentKey;

            var children = ChildrenSelector(item);
            if (children is { Count: > 0 })
            {
                BuildParentLookupRecursive(children, key, lookup);
            }
        }
    }
}
