XFEToolBox

using System.IO;
using System.Runtime.InteropServices;
using System.Windows;
using System.Windows.Interop;
using System.Windows.Media;
using System.Windows.Media.Imaging;

namespace XFEToolBox.Client.Utilities;

/// <summary>
/// Shell extensions may perform slow disk/COM work. Never run them on a dispatcher or a search worker.
/// Only realized result bindings ask for icons; obsolete requests can be cancelled without throwing.
/// </summary>
public sealed class ShellFileIconService : IDisposable
{
    public static ShellFileIconService Shared { get; } = new();

    private readonly object gate = new();
    private readonly int queueCapacity;
    private readonly int cacheCapacity;
    private readonly LinkedList<Work> queue = [];
    private readonly Dictionary<string, Work> pending = new(StringComparer.OrdinalIgnoreCase);
    private readonly Dictionary<string, LinkedListNode<CacheEntry>> cache = new(StringComparer.OrdinalIgnoreCase);
    private readonly LinkedList<CacheEntry> cacheOrder = [];
    private readonly Func<string, bool, bool, ImageSource?> iconLoader;
    private bool stopping;
    private long nextWaiterId;

    public ShellFileIconService() : this(2, 128, 2048) { }

    internal ShellFileIconService(int workerCount, int queueCapacity, int cacheCapacity,
        Func<string, bool, bool, ImageSource?>? iconLoader = null)
    {
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(workerCount);
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(queueCapacity);
        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(cacheCapacity);
        this.queueCapacity = queueCapacity;
        this.cacheCapacity = cacheCapacity;
        this.iconLoader = iconLoader ?? LoadShellIcon;
        for (var i = 0; i < workerCount; i++)
        {
            var thread = new Thread(WorkerLoop)
            {
                IsBackground = true,
                Name = $"XFEToolBox Shell icons {i + 1}",
                Priority = ThreadPriority.BelowNormal
            };
            thread.SetApartmentState(ApartmentState.STA);
            thread.Start();
        }
    }

    /// <summary>Returns a frozen cross-thread bitmap, or null when unavailable/cancelled/overloaded.</summary>
    public Task<ImageSource?> GetIconAsync(string path, bool isDirectory, CancellationToken cancellationToken = default)
    {
        if (cancellationToken.IsCancellationRequested) return Task.FromResult<ImageSource?>(null);
        var descriptor = Describe(path, isDirectory);
        Waiter waiter;
        Work work;
        lock (gate)
        {
            if (stopping) return Task.FromResult<ImageSource?>(null);
            if (cache.TryGetValue(descriptor.Key, out var cached))
            {
                if (cached.Value.ExpiresAt > DateTime.UtcNow)
                {
                    cacheOrder.Remove(cached);
                    cacheOrder.AddFirst(cached);
                    return Task.FromResult(cached.Value.Image);
                }
                cache.Remove(descriptor.Key);
                cacheOrder.Remove(cached);
            }
            if (!pending.TryGetValue(descriptor.Key, out work!))
            {
                // Prefer current visible results to stale queued results when typing/scrolling rapidly.
                if (queue.Count >= queueCapacity && queue.First is { } oldest) DropQueued(oldest.Value);
                work = new Work(descriptor);
                work.Node = queue.AddLast(work);
                pending.Add(descriptor.Key, work);
            }
            waiter = new Waiter(++nextWaiterId);
            work.Waiters.Add(waiter.Id, waiter);
            Monitor.Pulse(gate);
        }
        if (cancellationToken.CanBeCanceled)
        {
            var registration = cancellationToken.Register(() => CancelWaiter(work, waiter));
            lock (gate)
            {
                // Cancellation can run synchronously during Register; do not leak that registration.
                if (!waiter.Completion.Task.IsCompleted) waiter.Registration = registration;
                else registration.Unregister();
            }
        }
        return waiter.Completion.Task;
    }

    private void CancelWaiter(Work work, Waiter waiter)
    {
        lock (gate)
        {
            work.Waiters.Remove(waiter.Id);
            waiter.Registration.Unregister();
            waiter.Completion.TrySetResult(null);
            if (work.Waiters.Count == 0 && work.Node is not null) DropQueued(work);
        }
    }

    private void DropQueued(Work work)
    {
        if (work.Node is not null) queue.Remove(work.Node);
        work.Node = null;
        pending.Remove(work.Descriptor.Key);
        foreach (var waiter in work.Waiters.Values)
        {
            waiter.Registration.Unregister();
            waiter.Completion.TrySetResult(null);
        }
        work.Waiters.Clear();
    }

    private void WorkerLoop()
    {
        var comResult = OleInitialize(IntPtr.Zero);
        try
        {
            while (true)
            {
                Work work;
                lock (gate)
                {
                    while (!stopping && queue.Count == 0) Monitor.Wait(gate);
                    if (stopping) return;
                    work = queue.Last!.Value;
                    queue.RemoveLast();
                    work.Node = null;
                }
                ImageSource? image = null;
                try
                {
                    image = iconLoader(work.Descriptor.ShellPath, work.Descriptor.IsDirectory, work.Descriptor.UseAttributes);
                    if (image is not null && !image.IsFrozen)
                    {
                        if (image.CanFreeze) image.Freeze();
                        else image = null;
                    }
                }
                catch (Exception) { /* Broken/missing files and Shell handlers retain the lightweight glyph. */ }
                lock (gate)
                {
                    pending.Remove(work.Descriptor.Key);
                    if (!stopping)
                    {
                        var expiry = DateTime.UtcNow.Add(image is null ? TimeSpan.FromSeconds(10) :
                            work.Descriptor.UseAttributes ? TimeSpan.FromMinutes(10) : TimeSpan.FromMinutes(2));
                        var node = cacheOrder.AddFirst(new CacheEntry(work.Descriptor.Key, image, expiry));
                        cache[work.Descriptor.Key] = node;
                        while (cache.Count > cacheCapacity && cacheOrder.Last is { } last)
                        {
                            cache.Remove(last.Value.Key);
                            cacheOrder.RemoveLast();
                        }
                    }
                    foreach (var waiter in work.Waiters.Values)
                    {
                        waiter.Registration.Unregister();
                        waiter.Completion.TrySetResult(stopping ? null : image);
                    }
                    work.Waiters.Clear();
                }
            }
        }
        finally { if (comResult >= 0) OleUninitialize(); }
    }

    public void Dispose()
    {
        lock (gate)
        {
            if (stopping) return;
            stopping = true;
            while (queue.First is { } first) DropQueued(first.Value);
            // Do not join a possibly slow third-party Shell handler on a UI dispatcher.
            foreach (var work in pending.Values)
                foreach (var waiter in work.Waiters.Values)
                {
                    waiter.Registration.Unregister();
                    waiter.Completion.TrySetResult(null);
                }
            cache.Clear();
            cacheOrder.Clear();
            Monitor.PulseAll(gate);
        }
    }

    internal (int Queued, int Pending, int Cached) Counts
    {
        get { lock (gate) return (queue.Count, pending.Count, cache.Count); }
    }

    private static Descriptor Describe(string path, bool isDirectory)
    {
        if (isDirectory) return new("directory", "folder", true, true);
        var extension = Path.GetExtension(path);
        if (extension.Equals(".exe", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".lnk", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".ico", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".url", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".dll", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".cpl", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".scr", StringComparison.OrdinalIgnoreCase) ||
            extension.Equals(".appref-ms", StringComparison.OrdinalIgnoreCase))
            return new("path:" + path, path, false, false);
        // Attribute-only queries deliberately do not stat individual results or open document contents.
        return new("type:" + extension, extension.Length == 0 ? "file" : "file" + extension, false, true);
    }

    private static ImageSource? LoadShellIcon(string path, bool isDirectory, bool useAttributes)
    {
        var flags = 0x000000100u; // SHGFI_ICON | SHGFI_LARGEICON (32px, not the tiny 16px icon).
        if (useAttributes) flags |= 0x000000010; // SHGFI_USEFILEATTRIBUTES
        else flags |= 0x000000020; // SHGFI_ADDOVERLAYS, including shortcut overlays.
        var result = SHGetFileInfo(path, isDirectory ? 0x10u : 0x80u, out var info,
            (uint)Marshal.SizeOf<ShellFileInfo>(), flags);
        if (info.Icon == IntPtr.Zero) return null;
        try
        {
            if (result == UIntPtr.Zero) return null;
            var bitmap = Imaging.CreateBitmapSourceFromHIcon(info.Icon, Int32Rect.Empty,
                BitmapSizeOptions.FromEmptyOptions());
            bitmap.Freeze();
            return bitmap;
        }
        finally { DestroyIcon(info.Icon); }
    }

    private sealed record Descriptor(string Key, string ShellPath, bool IsDirectory, bool UseAttributes);
    private sealed record CacheEntry(string Key, ImageSource? Image, DateTime ExpiresAt);
    private sealed class Work(Descriptor descriptor)
    {
        public Descriptor Descriptor { get; } = descriptor;
        public LinkedListNode<Work>? Node { get; set; }
        public Dictionary<long, Waiter> Waiters { get; } = [];
    }
    private sealed class Waiter(long id)
    {
        public long Id { get; } = id;
        public TaskCompletionSource<ImageSource?> Completion { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
        public CancellationTokenRegistration Registration { get; set; }
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    private struct ShellFileInfo
    {
        public IntPtr Icon;
        public int IconIndex;
        public uint Attributes;
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)] public string DisplayName;
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 80)] public string TypeName;
    }
    [DllImport("shell32.dll", EntryPoint = "SHGetFileInfoW", CharSet = CharSet.Unicode)]
    private static extern UIntPtr SHGetFileInfo(string path, uint attributes, out ShellFileInfo info, uint size, uint flags);
    [DllImport("user32.dll")]
    [return: MarshalAs(UnmanagedType.Bool)]
    private static extern bool DestroyIcon(IntPtr icon);
    [DllImport("ole32.dll")]
    private static extern int OleInitialize(IntPtr reserved);
    [DllImport("ole32.dll")]
    private static extern void OleUninitialize();
}