@* 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 : "");

        // Checkbox-mode tri-state computed for this node (folder = aggregate of
        // descendant leaves; leaf = present-in-SelectedKeys).
        var checkState = SelectionMode == TreeViewSelectionMode.Checkbox
            ? ComputeCheckState(item)
            : CheckState.Unchecked;

        <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)" @oncontextmenu:stopPropagation="@(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>
                }
                @if (SelectionMode == TreeViewSelectionMode.Checkbox)
                {
                    <input type="checkbox"
                           class="form-check-input tv-checkbox @(checkState == CheckState.Indeterminate ? "tv-checkbox-indeterminate" : "")"
                           @ref="_checkboxRefs[KeyStr(key)]"
                           checked="@(checkState == CheckState.Checked)"
                           @onchange="() => OnCheckboxToggle(item)"
                           @onclick:stopPropagation />
                }
                <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; }

    // ── Checkbox-selection mode (additive; SelectionMode=Single keeps prior behaviour) ──
    [Parameter] public TreeViewSelectionMode SelectionMode { get; set; } = TreeViewSelectionMode.Single;
    [Parameter] public HashSet<object>? SelectedKeys { get; set; }
    [Parameter] public EventCallback<HashSet<object>> SelectedKeysChanged { get; set; }

    private readonly Dictionary<string, ElementReference> _checkboxRefs = new();

    private enum CheckState { Unchecked, Checked, Indeterminate }

    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;
            // The context-menu element may have been removed (menu dismissed
            // during render) or the circuit disconnected — both are expected.
            try { await _contextMenuRef.FocusAsync(); }
            catch (Microsoft.JSInterop.JSException) { }
            catch (Microsoft.JSInterop.JSDisconnectedException) { }
            catch (InvalidOperationException) { }
        }

        if (firstRender && StorageKey != null)
        {
            string? json = null;
            try
            {
                json = await JSRuntime.InvokeAsync<string?>("treeviewStorage.load", StorageKey);
            }
            catch (Microsoft.JSInterop.JSDisconnectedException)
            {
                // Circuit disconnected before the storage read completed — there
                // is nothing to restore and nothing to log.
            }
            _storageLoaded = true;

            // CentralUI-018: a corrupt or wrong-shaped treeviewStorage payload
            // must not throw out of OnAfterRenderAsync. Guard the deserialize
            // and treat an unparseable payload as "no prior state".
            List<string>? keys = null;
            if (json != null)
            {
                try
                {
                    keys = System.Text.Json.JsonSerializer.Deserialize<List<string>>(json);
                }
                catch (System.Text.Json.JsonException)
                {
                    keys = null;
                }
            }

            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 or a corrupt payload (no usable prior
                // state) — fall back to InitiallyExpanded.
                _initialExpansionApplied = true;
                ApplyInitialExpansion(_items);
            }

            StateHasChanged();
        }

        // Apply checkbox tri-state (`indeterminate`) after every render in
        // Checkbox mode. Blazor doesn't bind input.indeterminate natively.
        if (SelectionMode == TreeViewSelectionMode.Checkbox && _items is { Count: > 0 })
        {
            await ApplyIndeterminateStateAsync();
        }
    }

    private async Task ApplyIndeterminateStateAsync()
    {
        foreach (var (keyStr, elemRef) in _checkboxRefs)
        {
            var item = FindItemByKey(_items!, keyStr);
            if (item is null) continue;
            var state = ComputeCheckState(item);
            try
            {
                await JSRuntime.InvokeVoidAsync(
                    "treeviewStorage.setIndeterminate",
                    elemRef,
                    state == CheckState.Indeterminate);
            }
            catch (Microsoft.JSInterop.JSDisconnectedException) { /* circuit gone */ }
            catch (Microsoft.JSInterop.JSException) { /* element gone */ }
        }
    }

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

    // ── Checkbox-selection helpers ──────────────────────────────────────────
    // Folder = aggregate of its descendant LEAVES (we don't track folder keys
    // in SelectedKeys — only leaf keys are persisted). A folder is Checked when
    // every descendant leaf is in SelectedKeys, Unchecked when none are, and
    // Indeterminate otherwise.
    private CheckState ComputeCheckState(TItem item)
    {
        var children = ChildrenSelector(item);
        var isBranch = HasChildrenSelector(item);
        if (!isBranch || children is null || children.Count == 0)
        {
            var leafKey = KeySelector(item);
            return SelectedKeys != null && SelectedKeys.Contains(leafKey)
                ? CheckState.Checked
                : CheckState.Unchecked;
        }

        var total = 0;
        var selected = 0;
        CountDescendantLeaves(item, ref total, ref selected);
        if (total == 0) return CheckState.Unchecked;
        if (selected == 0) return CheckState.Unchecked;
        if (selected == total) return CheckState.Checked;
        return CheckState.Indeterminate;
    }

    private void CountDescendantLeaves(TItem item, ref int total, ref int selected)
    {
        var children = ChildrenSelector(item);
        var hasChildren = HasChildrenSelector(item) && children is { Count: > 0 };
        if (!hasChildren)
        {
            total++;
            if (SelectedKeys != null && SelectedKeys.Contains(KeySelector(item)))
            {
                selected++;
            }
            return;
        }

        foreach (var child in children!)
        {
            CountDescendantLeaves(child, ref total, ref selected);
        }
    }

    private void CollectDescendantLeafKeys(TItem item, List<object> sink)
    {
        var children = ChildrenSelector(item);
        var hasChildren = HasChildrenSelector(item) && children is { Count: > 0 };
        if (!hasChildren)
        {
            sink.Add(KeySelector(item));
            return;
        }

        foreach (var child in children!)
        {
            CollectDescendantLeafKeys(child, sink);
        }
    }

    private async Task OnCheckboxToggle(TItem item)
    {
        // SelectedKeys is the source of truth — copy-on-write so consumers
        // observe a fresh reference and Blazor reliably re-renders.
        var current = SelectedKeys is null
            ? new HashSet<object>()
            : new HashSet<object>(SelectedKeys);

        var leaves = new List<object>();
        CollectDescendantLeafKeys(item, leaves);
        if (leaves.Count == 0) return;

        // Folder-toggle semantics: if every descendant leaf is currently selected,
        // uncheck them all; otherwise select all. Leaf nodes have leaves = { self }
        // so this simplifies to a plain toggle.
        var allSelected = leaves.All(current.Contains);
        if (allSelected)
        {
            foreach (var k in leaves) current.Remove(k);
        }
        else
        {
            foreach (var k in leaves) current.Add(k);
        }

        SelectedKeys = current;
        await SelectedKeysChanged.InvokeAsync(current);
    }

    private TItem? FindItemByKey(IReadOnlyList<TItem> items, string keyStr)
    {
        foreach (var item in items)
        {
            if (KeyStr(KeySelector(item)) == keyStr) return item;
            var children = ChildrenSelector(item);
            if (children is { Count: > 0 })
            {
                var found = FindItemByKey(children, keyStr);
                if (found is not null) return found;
            }
        }
        return default;
    }
}