DodoSSH/src/DodoSSH.Client.App/Views/VirtualizingWrapPanel.cs

using System.Collections.Specialized;
using Avalonia;
using Avalonia.Controls;
using Avalonia.Input;
using Avalonia.Layout;

namespace DodoSSH.Client.App.Views;

/// <summary>
/// A wrap panel that builds only the items near the part of it on screen.
/// </summary>
/// <remarks>
/// <para>
/// The hosts board is a list per section, each laid out in rows of cards. It was a plain
/// <see cref="WrapPanel"/>, which builds, measures and draws every card on the board whether it is in view or
/// not — some twenty-five controls a card, so a keychain of three hundred hosts took seconds to open and to
/// redraw after a filter was cleared, and a thousand took several times that. Avalonia has no virtualising
/// wrap panel of its own; this is one, for items of a fixed width.
/// </para>
/// <para>
/// <b>Widths are the first item's; heights are per row.</b> A card's width is set by its style, so the
/// number of columns is arithmetic. Its height is not quite — a card wearing many tags grows a line — so each
/// row's height is the tallest card built in it, and a row not yet built is assumed to be as tall as the last
/// one measured. When a row turns out different the rows below it move, which is the ordinary price of
/// estimating, and is only ever seen while scrolling into rows nobody has looked at yet.
/// </para>
/// <para>
/// <b>What is on screen is read from the effective viewport</b>, which is the part of this panel the
/// scroller around the whole board shows — the section's own list never scrolls, so its own scroller has
/// nothing to say. A section scrolled out of view builds nothing at all past the one card measured for size.
/// </para>
/// <para>
/// Everything else a list does is left to the list. Containers come from its generator and are prepared by
/// it, so selection, the card template, the context menu and the band all work on the containers that exist.
/// A container holding the keyboard is never put away while it does, so arrowing past the edge of the view
/// does not drop the focus on the floor.
/// </para>
/// </remarks>
internal sealed class VirtualizingWrapPanel : VirtualizingPanel
{
    /// <summary>How far past the visible part, above and below, items are built ahead of being scrolled to.</summary>
    private const double Overscan = 240;

    /// <summary>What a row is assumed to be before anything in the panel has been measured.</summary>
    private const double FallbackRowHeight = 124;

    private readonly Dictionary<int, Control> realized = [];

    private readonly Dictionary<object, Stack<Control>> recycled = [];

    /// <summary>The pool each container this panel made goes back to; items that are their own containers are not in it.</summary>
    private readonly Dictionary<Control, object> poolOf = [];

    private readonly Dictionary<int, double> rowHeights = [];

    private Rect viewport;

    private bool viewportKnown;

    private Size itemSize;

    private int columns = 1;

    /// <summary>An item asked to be scrolled to, kept built until a measure has seen it in range.</summary>
    private int? pinned;

    public VirtualizingWrapPanel()
    {
        EffectiveViewportChanged += (_, e) =>
        {
            if (viewportKnown && e.EffectiveViewport == viewport)
            {
                return;
            }

            viewport = e.EffectiveViewport;
            viewportKnown = true;
            InvalidateMeasure();
        };
    }

    /// <inheritdoc />
    protected override Size MeasureOverride(Size availableSize)
    {
        var items = Items;

        if (items.Count == 0)
        {
            RecycleAllExcept(static _ => false);
            return default;
        }

        if (itemSize == default)
        {
            var first = RealizedOrNew(0);
            first.Measure(Size.Infinity);
            itemSize = first.DesiredSize;
        }

        var width = double.IsInfinity(availableSize.Width)
            ? itemSize.Width * items.Count
            : availableSize.Width;

        var nowColumns = Math.Max(1, (int)Math.Floor(width / Math.Max(1, itemSize.Width)));

        if (nowColumns != columns)
        {
            columns = nowColumns;
            rowHeights.Clear();
        }

        var tops = RowTops(items.Count);
        var (firstIndex, lastIndex) = VisibleRange(tops, items.Count);

        RecycleAllExcept(index =>
            (index >= firstIndex && index <= lastIndex)
            || index == pinned
            || realized[index].IsKeyboardFocusWithin);

        for (var index = firstIndex; index <= lastIndex; index++)
        {
            var container = RealizedOrNew(index);
            container.Measure(new Size(itemSize.Width, double.PositiveInfinity));

            var row = index / columns;
            rowHeights[row] = Math.Max(rowHeights.GetValueOrDefault(row), container.DesiredSize.Height);
        }

        if (pinned is { } kept && kept >= firstIndex && kept <= lastIndex)
        {
            pinned = null;
        }

        tops = RowTops(items.Count);

        return new Size(Math.Min(width, columns * itemSize.Width), tops[^1]);
    }

    /// <inheritdoc />
    protected override Size ArrangeOverride(Size finalSize)
    {
        var tops = RowTops(Items.Count);

        foreach (var (index, container) in realized)
        {
            container.Arrange(ItemRect(index, tops));
        }

        return finalSize;
    }

    /// <inheritdoc />
    protected override void OnItemsChanged(IReadOnlyList<object?> items, NotifyCollectionChangedEventArgs e)
    {
        base.OnItemsChanged(items, e);

        // Every change is answered as a reset. The board swaps its lists in whole, so an incremental answer
        // would be code for a case that does not happen; and indexes, row heights and the pinned item all
        // describe a list that is no longer the one being shown.
        RecycleAllExcept(static _ => false);
        rowHeights.Clear();
        pinned = null;
        InvalidateMeasure();
    }

    /// <inheritdoc />
    protected override Control? ScrollIntoView(int index)
    {
        if (index < 0 || index >= Items.Count)
        {
            return null;
        }

        var container = RealizedOrNew(index);
        container.Measure(new Size(itemSize.Width, double.PositiveInfinity));
        container.Arrange(ItemRect(index, RowTops(Items.Count)));

        pinned = index;
        container.BringIntoView();

        return container;
    }

    /// <inheritdoc />
    protected override Control? ContainerFromIndex(int index) => realized.GetValueOrDefault(index);

    /// <inheritdoc />
    protected override int IndexFromContainer(Control container)
    {
        foreach (var (index, held) in realized)
        {
            if (ReferenceEquals(held, container))
            {
                return index;
            }
        }

        return -1;
    }

    /// <inheritdoc />
    protected override IEnumerable<Control>? GetRealizedContainers() => realized.Values;

    /// <inheritdoc />
    /// <remarks>
    /// Left and right step one card, up and down one row of them — the movements the plain wrap panel's own
    /// navigation offered, so the board's arrows do what they did, with the card stepped to built on demand.
    /// </remarks>
    protected override IInputElement? GetControl(NavigationDirection direction, IInputElement? from, bool wrap)
    {
        var count = Items.Count;

        if (count == 0)
        {
            return null;
        }

        var current = from is Control control ? IndexFromContainer(control) : -1;

        var target = direction switch
        {
            NavigationDirection.First => 0,
            NavigationDirection.Last => count - 1,
            NavigationDirection.Left or NavigationDirection.Previous => current - 1,
            NavigationDirection.Right or NavigationDirection.Next => current + 1,
            NavigationDirection.Up => current - columns,
            NavigationDirection.Down => current + columns,
            _ => -1,
        };

        if (current < 0 && direction is not (NavigationDirection.First or NavigationDirection.Last))
        {
            target = 0;
        }

        if (target < 0 || target >= count)
        {
            if (!wrap)
            {
                return null;
            }

            target = (target % count + count) % count;
        }

        return ScrollIntoView(target);
    }

    /// <summary>The top of every row, and one more entry past the last row for the total height.</summary>
    private double[] RowTops(int count)
    {
        var rows = (count + columns - 1) / columns;
        var tops = new double[rows + 1];
        var estimate = itemSize.Height > 0 ? itemSize.Height : FallbackRowHeight;

        for (var row = 0; row < rows; row++)
        {
            tops[row + 1] = tops[row] + rowHeights.GetValueOrDefault(row, estimate);
        }

        return tops;
    }

    /// <summary>The first and last item in the rows that overlap the visible part, with the overscan.</summary>
    private (int First, int Last) VisibleRange(double[] tops, int count)
    {
        // Before the first viewport arrives, a screenful from the top: the panel has to measure to something.
        var top = viewportKnown ? viewport.Top - Overscan : 0;
        var bottom = viewportKnown ? viewport.Bottom + Overscan : 1200;

        var rows = tops.Length - 1;
        var firstRow = -1;
        var lastRow = -1;

        for (var row = 0; row < rows; row++)
        {
            if (tops[row + 1] < top)
            {
                continue;
            }

            if (tops[row] > bottom)
            {
                break;
            }

            if (firstRow < 0)
            {
                firstRow = row;
            }

            lastRow = row;
        }

        return firstRow < 0
            ? (0, -1)
            : (firstRow * columns, Math.Min(count - 1, ((lastRow + 1) * columns) - 1));
    }

    private Rect ItemRect(int index, double[] tops)
    {
        var row = index / columns;
        var column = index % columns;
        var height = row + 1 < tops.Length ? tops[row + 1] - tops[row] : itemSize.Height;

        return new Rect(column * itemSize.Width, tops[Math.Min(row, tops.Length - 1)], itemSize.Width, height);
    }

    /// <summary>The container for an item, built or brought back from the pool if it is not already up.</summary>
    private Control RealizedOrNew(int index)
    {
        if (realized.TryGetValue(index, out var existing))
        {
            return existing;
        }

        var generator = ItemContainerGenerator!;
        var item = Items[index];
        Control container;

        if (generator.NeedsContainer(item, index, out var recycleKey))
        {
            var key = recycleKey ?? typeof(object);

            if (recycled.TryGetValue(key, out var pool) && pool.TryPop(out var reused))
            {
                container = reused;
                container.IsVisible = true;
            }
            else
            {
                container = generator.CreateContainer(item, index, recycleKey);
                poolOf[container] = key;
                AddInternalChild(container);
            }

            generator.PrepareItemContainer(container, item, index);
            realized[index] = container;
            generator.ItemContainerPrepared(container, item, index);
        }
        else
        {
            // An item that is its own container: built by whoever made the item, and never pooled.
            container = (Control)item!;
            AddInternalChild(container);
            realized[index] = container;
        }

        return container;
    }

    /// <summary>Puts away every built container the predicate does not keep.</summary>
    private void RecycleAllExcept(Func<int, bool> keep)
    {
        foreach (var index in realized.Keys.ToList())
        {
            if (keep(index))
            {
                continue;
            }

            var container = realized[index];
            realized.Remove(index);

            if (poolOf.TryGetValue(container, out var key))
            {
                ItemContainerGenerator!.ClearItemContainer(container);
                container.IsVisible = false;

                if (!recycled.TryGetValue(key, out var pool))
                {
                    recycled[key] = pool = new Stack<Control>();
                }

                pool.Push(container);
            }
            else
            {
                RemoveInternalChild(container);
            }
        }
    }
}