XFEEDToolBox

using System.Runtime.CompilerServices;
using XFEEDToolBox.Core.Models;

namespace XFEEDToolBox.Core.Galaxy;

public sealed record GalaxyCellStreamingOptions
{
    public int BatchSize { get; init; } = 128;
    public int MaximumCachedCells { get; init; } = 2048;
    public int MaximumCachedSystems { get; init; } = 200000;
}

public sealed record GalaxyCellStreamUpdate(IReadOnlyList<GalacticSystem> Systems, int QueriedBatches,
    int CoveredCells, int PendingCells, int UnverifiedCells, int RefinedBatches, bool Finished,
    bool CoverageComplete, bool FromCache, long OmittedCandidates);

/// <summary>
/// Streams original mass cells. A complete receipt covers only its exact requested
/// cells; neither a coarse mass code nor a small row count proves wider coverage.
/// </summary>
public sealed class GalaxyCellStreamingSession : IDisposable
{
    private readonly record struct Key(Coordinates Origin, ulong Address);
    private sealed record Entry(IReadOnlyList<GalacticSystem> Rows, bool Complete, long Used);
    private sealed record Work(Coordinates Origin, ulong[] Addresses, CancellationTokenSource Stop);
    private readonly GalaxyCellStreamingOptions options;
    private readonly Dictionary<Key, Entry> cache = [];
    private readonly SemaphoreSlim gate = new(1, 1);
    private readonly object sync = new();
    private CancellationTokenSource? active;
    private Work? work;
    private long generation, clock, epoch;
    private bool disposed;
    public int CachedCells { get { lock (sync) return cache.Count; } }

    public GalaxyCellStreamingSession(GalaxyCellStreamingOptions? options = null)
    {
        this.options = options ?? new();
        if (this.options.BatchSize is < 1 or > 512 || this.options.MaximumCachedCells < 1 || this.options.MaximumCachedSystems < 1)
            throw new ArgumentException("原星区缓存配置无效。");
    }

    public async IAsyncEnumerable<GalaxyCellStreamUpdate> StreamAsync(GalaxyCellQueryPlan plan,
        Func<Coordinates, IReadOnlyList<double>, IReadOnlyList<ulong>, CancellationToken, Task<GalaxyRegionResult>> query,
        [EnumeratorCancellation] CancellationToken token = default)
    {
        ArgumentNullException.ThrowIfNull(plan); ArgumentNullException.ThrowIfNull(query);
        var version = Interlocked.Increment(ref generation);
        using var stop = CancellationTokenSource.CreateLinkedTokenSource(token);
        var wanted = plan.Cells.Select(c => c.Address).ToHashSet();
        lock (sync)
        {
            ObjectDisposedException.ThrowIf(disposed, this);
            if (work is { } current && (current.Origin != plan.AbsoluteOrigin || !current.Addresses.Any(wanted.Contains))) Cancel(current.Stop);
            Cancel(active); active = stop;
        }
        await gate.WaitAsync(stop.Token);
        try
        {
            var seen = new HashSet<long>(); var covered = new HashSet<ulong>(); var unverified = new HashSet<ulong>();
            var pending = new Queue<ulong[]>(); var missing = new List<ulong>(); var cachedRows = new List<GalacticSystem>();
            var queried = 0; var refined = 0;
            void Current()
            {
                stop.Token.ThrowIfCancellationRequested();
                if (version != Volatile.Read(ref generation)) throw new OperationCanceledException(stop.Token);
            }
            IReadOnlyList<GalacticSystem> NewRows(IEnumerable<GalacticSystem> rows) =>
                rows.Where(s => s.SystemAddress is { } id && seen.Add(id)).ToArray();
            GalaxyCellStreamUpdate Update(IReadOnlyList<GalacticSystem> rows, bool fromCache, bool finished = false) =>
                new(rows, queried, covered.Count, pending.Sum(b => b.Length), unverified.Count, refined, finished,
                    finished && !plan.Truncated && unverified.Count == 0 && covered.Count == wanted.Count,
                    fromCache, plan.OmittedCandidates);
            Current();
            lock (sync)
                foreach (var cell in plan.Cells)
                {
                    var key = new Key(plan.AbsoluteOrigin, cell.Address);
                    if (cache.TryGetValue(key, out var entry))
                    {
                        cache[key] = entry with { Used = ++clock }; cachedRows.AddRange(entry.Rows);
                        if (entry.Complete) { covered.Add(cell.Address); continue; }
                    }
                    missing.Add(cell.Address);
                }
            foreach (var batch in missing.Chunk(options.BatchSize)) pending.Enqueue(batch);
            if (cachedRows.Count > 0 || covered.Count > 0) yield return Update(NewRows(cachedRows), true);
            while (pending.TryDequeue(out var addresses))
            {
                Current();
                using var queryStop = new CancellationTokenSource();
                using var callerCancellation = token.Register(() => Cancel(queryStop));
                var current = new Work(plan.AbsoluteOrigin, addresses, queryStop);
                long acceptedEpoch;
                lock (sync) { Current(); ObjectDisposedException.ThrowIf(disposed, this); work = current; acceptedEpoch = epoch; }
                GalaxyRegionResult result;
                try { result = await query(plan.Descriptor.Centre, plan.Descriptor.MassRadii, addresses, queryStop.Token); }
                finally { lock (sync) if (ReferenceEquals(work, current)) work = null; }
                queried++;
                if (result.Origin != plan.AbsoluteOrigin) throw new InvalidDataException("原星区分页的银河坐标原点缺失或已变化。");
                var groups = ValidateRows(result.Systems, addresses);
                var complete = CoversRequestedCells(result, addresses.Length) && result.MinimumMassCode == (int)addresses.Min(a => a & 7);
                // A useful accepted query may finish after a pan. Cache its real
                // cells before rejecting publication into the superseded view.
                lock (sync)
                {
                    if (acceptedEpoch == epoch && !disposed)
                    {
                        foreach (var address in addresses)
                        {
                            var key = new Key(plan.AbsoluteOrigin, address);
                            var rows = groups.GetValueOrDefault(address) ?? [];
                            if (!complete && cache.TryGetValue(key, out var old))
                                rows = old.Rows.Concat(rows).DistinctBy(s => s.SystemAddress).ToArray();
                            cache[key] = new(rows, complete, ++clock);
                        }
                        Trim();
                    }
                }
                Current();
                if (complete) foreach (var address in addresses) covered.Add(address);
                else if (addresses.Length > 1)
                {
                    // Split the ID list, never its spatial bounds: full edge-cell
                    // rows must survive retries of a capped native batch.
                    var half = addresses.Length / 2;
                    pending.Enqueue(addresses[..half]); pending.Enqueue(addresses[half..]); refined++;
                }
                else unverified.Add(addresses[0]);
                yield return Update(NewRows(result.Systems), false);
            }
            Current();
            yield return Update([], false, true);
        }
        finally
        {
            lock (sync) if (ReferenceEquals(active, stop)) active = null;
            gate.Release();
        }
    }

    public static bool CoversRequestedCells(GalaxyRegionResult result, int count) =>
        count > 0 && result.Scope == "requested-cells" && result.RequestedCells == count && result.Complete &&
        result.HasCoverageEvidence && result.Error.Length == 0 && !result.Truncated && result.WithinReturnedRowBudget &&
        result.FailedCells == 0 && result.ProcessedCells == count && result.TotalCells == count;

    private static Dictionary<ulong, IReadOnlyList<GalacticSystem>> ValidateRows(IReadOnlyList<GalacticSystem> rows, ulong[] addresses)
    {
        var requested = addresses.ToHashSet();
        var groups = new Dictionary<ulong, List<GalacticSystem>>();
        var seen = new Dictionary<long, GalacticSystem>();
        foreach (var row in rows)
        {
            if (row.SystemAddress is not { } id || id <= 0 || id >= 1L << 55 || !row.Position.IsValid)
                throw new InvalidDataException("原星区分页返回了无效的恒星地址或坐标。");
            var cell = GalaxyCellQueryPlan.CellAddress((ulong)id);
            if (!requested.Contains(cell)) throw new InvalidDataException("原星区分页返回了请求之外的恒星。");
            if (seen.TryGetValue(id, out var previous))
            {
                if (previous.Position != row.Position || previous.Name != row.Name)
                    throw new InvalidDataException("原星区分页返回了相互冲突的恒星记录。");
                continue;
            }
            seen.Add(id, row);
            if (!groups.TryGetValue(cell, out var group)) groups[cell] = group = [];
            group.Add(row);
        }
        return groups.ToDictionary(p => p.Key, p => (IReadOnlyList<GalacticSystem>)p.Value.ToArray());
    }

    private void Trim()
    {
        var rows = cache.Values.Sum(e => e.Rows.Count);
        foreach (var pair in cache.OrderBy(p => p.Value.Used).ToArray())
        {
            if (cache.Count <= options.MaximumCachedCells && rows <= options.MaximumCachedSystems) break;
            cache.Remove(pair.Key); rows -= pair.Value.Rows.Count;
        }
    }
    public void Clear()
    {
        lock (sync) { epoch++; cache.Clear(); Cancel(active); if (work is { } current) Cancel(current.Stop); }
    }
    public void Dispose()
    {
        lock (sync) { if (disposed) return; disposed = true; Clear(); }
    }
    private static void Cancel(CancellationTokenSource? source)
    { try { source?.Cancel(); } catch (ObjectDisposedException) { } }
}