av-parser/src/AvParser.Core/Proxies/ProxyPool.cs

using AvParser.Core.Proxies.Selection;

namespace AvParser.Core.Proxies;

/// <inheritdoc cref="IProxyPool" />
public sealed class ProxyPool : IProxyPool
{
    private readonly IReadOnlyList<IProxySource> _sources;
    private readonly IProxyProbe _probe;
    private readonly TimeProvider _time;
    private readonly Random? _random;
    private readonly Lock _gate = new();

    private readonly Dictionary<string, ProxyEntry> _byKey = new(StringComparer.Ordinal);
    private List<ProxyEntry> _entries = [];
    private IProxySelectionStrategy _strategy;
    private ProxyOptions _options;

    /// <summary>Creates a pool over the given sources.</summary>
    /// <param name="sources">Every list the app knows about; order decides nothing.</param>
    /// <param name="probe">Liveness checker.</param>
    /// <param name="options">Initial options.</param>
    /// <param name="timeProvider">Clock; tests inject a fake one to exercise quarantine expiry.</param>
    /// <param name="random">Randomness for weighted selection; tests seed it.</param>
    public ProxyPool(
        IEnumerable<IProxySource> sources,
        IProxyProbe probe,
        ProxyOptions? options = null,
        TimeProvider? timeProvider = null,
        Random? random = null
    )
    {
        ArgumentNullException.ThrowIfNull(sources);

        _sources = sources.ToArray();
        _probe = probe ?? throw new ArgumentNullException(nameof(probe));
        _time = timeProvider ?? TimeProvider.System;
        _random = random;
        _options = (options ?? new ProxyOptions()).Validated();
        _strategy = ProxySelectionStrategyFactory.Create(_options.Rotation, random);
    }

    /// <inheritdoc />
    public IReadOnlyList<ProxyEntry> Entries
    {
        get
        {
            lock (_gate)
            {
                return _entries;
            }
        }
    }

    /// <inheritdoc />
    public ProxyOptions Options
    {
        get
        {
            lock (_gate)
            {
                return _options;
            }
        }
    }

    /// <inheritdoc />
    public int LiveCount
    {
        get
        {
            var now = _time.GetUtcNow();

            lock (_gate)
            {
                return _entries.Count(entry => entry.Health == ProxyHealthState.Alive && entry.IsAvailable(now));
            }
        }
    }

    /// <inheritdoc />
    public event EventHandler? Changed;

    /// <inheritdoc />
    public void Configure(ProxyOptions options)
    {
        ArgumentNullException.ThrowIfNull(options);

        var validated = options.Validated();

        lock (_gate)
        {
            var rotationChanged = validated.Rotation != _options.Rotation;
            _options = validated;

            if (rotationChanged)
            {
                _strategy = ProxySelectionStrategyFactory.Create(validated.Rotation, _random);
            }
        }

        RaiseChanged();
    }

    /// <inheritdoc />
    public async Task<int> RefreshAsync(CancellationToken cancellationToken = default)
    {
        var options = Options;
        var collected = new List<(ProxyEndpoint Endpoint, ProxySourceKind Kind)>();

        foreach (var source in _sources)
        {
            if (source.Kind == ProxySourceKind.Feed && !options.UseFeed)
            {
                continue;
            }

            var proxies = await source.GetProxiesAsync(cancellationToken).ConfigureAwait(false);
            foreach (var endpoint in proxies)
            {
                collected.Add((endpoint, source.Kind));
            }
        }

        int count;
        lock (_gate)
        {
            var next = new List<ProxyEntry>(collected.Count);
            var seen = new HashSet<string>(StringComparer.Ordinal);

            foreach (var (endpoint, kind) in collected)
            {
                if (!Matches(endpoint, options) || !seen.Add(endpoint.Key))
                {
                    continue;
                }

                // Reuse the existing entry so a refresh does not wipe out everything the pool
                // has learned — free lists are republished every few minutes.
                if (_byKey.TryGetValue(endpoint.Key, out var existing))
                {
                    next.Add(existing);
                }
                else
                {
                    var created = new ProxyEntry(endpoint, kind);
                    _byKey[endpoint.Key] = created;
                    next.Add(created);
                }
            }

            foreach (var key in _byKey.Keys.Where(key => !seen.Contains(key)).ToArray())
            {
                _byKey.Remove(key);
            }

            _entries = next;
            _strategy.Reset();
            count = next.Count;
        }

        RaiseChanged();
        return count;
    }

    /// <inheritdoc />
    public async Task<ProxyLease?> AcquireAsync(CancellationToken cancellationToken = default)
    {
        var options = Options;

        if (options.HealthCheck == ProxyHealthCheck.Pool)
        {
            var entry = SelectAvailable();
            return entry is null ? null : new ProxyLease(this, entry);
        }

        // Lazy: verify the pick before handing it over, stepping past dead ones.
        for (var attempt = 0; attempt < options.LazyProbeAttempts; attempt++)
        {
            cancellationToken.ThrowIfCancellationRequested();

            var entry = SelectAvailable();
            if (entry is null)
            {
                return null;
            }

            if (entry.Health == ProxyHealthState.Alive)
            {
                return new ProxyLease(this, entry);
            }

            var result = await _probe.ProbeAsync(entry.Endpoint, options, cancellationToken).ConfigureAwait(false);
            var now = _time.GetUtcNow();

            if (result.Alive)
            {
                entry.RecordProbe(now, alive: true, result.Latency, error: null);
                RaiseChanged();
                return new ProxyLease(this, entry);
            }

            // Same reasoning as the sweep: this is a probe verdict, so it sidelines the entry
            // without inflating the request counters that weighted selection reads.
            entry.RecordProbe(now, alive: false, latency: null, result.Error, options.BaseQuarantine);

            lock (_gate)
            {
                _strategy.Report(entry, success: false);
            }

            RaiseChanged();
        }

        return null;
    }

    /// <inheritdoc />
    public async Task<int> SweepAsync(
        IProgress<ProxySweepProgress>? progress = null,
        CancellationToken cancellationToken = default
    )
    {
        var options = Options;
        var targets = Entries;
        var total = targets.Count;

        if (total == 0)
        {
            progress?.Report(new ProxySweepProgress(0, 0, 0));
            return 0;
        }

        using var limiter = new SemaphoreSlim(options.ProbeConcurrency, options.ProbeConcurrency);
        var checkedCount = 0;
        var aliveCount = 0;

        var work = targets.Select(async entry =>
        {
            await limiter.WaitAsync(cancellationToken).ConfigureAwait(false);
            try
            {
                var result = await _probe.ProbeAsync(entry.Endpoint, options, cancellationToken).ConfigureAwait(false);

                // A failed sweep probe sidelines the entry for one base window rather than
                // counting as a request failure: it is evidence about the proxy, not usage of it.
                entry.RecordProbe(
                    _time.GetUtcNow(),
                    result.Alive,
                    result.Latency,
                    result.Error,
                    quarantineOnFailure: options.BaseQuarantine
                );

                if (result.Alive)
                {
                    Interlocked.Increment(ref aliveCount);
                }
            }
            finally
            {
                limiter.Release();
                progress?.Report(
                    new ProxySweepProgress(
                        Interlocked.Increment(ref checkedCount),
                        total,
                        Volatile.Read(ref aliveCount)
                    )
                );
            }
        });

        await Task.WhenAll(work).ConfigureAwait(false);

        lock (_gate)
        {
            // A sweep can invalidate a sticky pick, so make the strategy re-choose.
            _strategy.Reset();
        }

        RaiseChanged();
        return aliveCount;
    }

    /// <inheritdoc />
    public async Task<int> WarmUpAsync(
        int targetLive,
        IProgress<ProxySweepProgress>? progress = null,
        CancellationToken cancellationToken = default
    )
    {
        ArgumentOutOfRangeException.ThrowIfLessThan(targetLive, 1);

        var options = Options;
        var candidates = WarmUpOrder();
        var total = candidates.Count;

        if (total == 0 || LiveCount >= targetLive)
        {
            progress?.Report(new ProxySweepProgress(0, 0, LiveCount));
            return LiveCount;
        }

        // Stops the remaining probes the moment the target is met. Linked so the caller's own
        // cancellation still wins.
        using var stop = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        using var limiter = new SemaphoreSlim(options.ProbeConcurrency, options.ProbeConcurrency);

        var checkedCount = 0;

        var work = candidates.Select(async entry =>
        {
            if (stop.IsCancellationRequested)
            {
                return;
            }

            try
            {
                await limiter.WaitAsync(stop.Token).ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
                return;
            }

            try
            {
                var result = await _probe.ProbeAsync(entry.Endpoint, options, stop.Token).ConfigureAwait(false);
                entry.RecordProbe(
                    _time.GetUtcNow(),
                    result.Alive,
                    result.Latency,
                    result.Error,
                    quarantineOnFailure: options.BaseQuarantine
                );
            }
            catch (OperationCanceledException)
            {
                // Either the target was reached or the caller gave up; neither is the proxy's fault.
                return;
            }
            finally
            {
                limiter.Release();
            }

            var live = LiveCount;
            progress?.Report(new ProxySweepProgress(Interlocked.Increment(ref checkedCount), total, live));

            if (live >= targetLive)
            {
                await stop.CancelAsync().ConfigureAwait(false);
            }
        });

        await Task.WhenAll(work).ConfigureAwait(false);

        cancellationToken.ThrowIfCancellationRequested();

        lock (_gate)
        {
            _strategy.Reset();
        }

        RaiseChanged();
        return LiveCount;
    }

    /// <summary>
    /// Orders candidates for a warm-up: what worked last time, then what looks most promising.
    /// </summary>
    /// <remarks>Public so the ordering can be asserted without running probes.</remarks>
    public IReadOnlyList<ProxyEntry> WarmUpOrder()
    {
        var now = _time.GetUtcNow();

        lock (_gate)
        {
            return
            [
                .. _entries
                    .Where(entry => entry.IsAvailable(now))
                    .OrderByDescending(entry => entry.Health == ProxyHealthState.Alive)
                    .ThenByDescending(entry => entry.WasAliveOnLastRun)
                    .ThenByDescending(entry => entry.SuccessCount > 0)
                    .ThenBy(entry => entry.Latency ?? TimeSpan.MaxValue)
                    .ThenByDescending(entry => entry.SuccessRate)
                    .ThenByDescending(entry => entry.Endpoint.Score),
            ];
        }
    }

    /// <summary>Records the outcome of a lease. Called by <see cref="ProxyLease"/>.</summary>
    internal void ReportOutcome(ProxyEntry entry, bool success, TimeSpan? latency, string? error)
    {
        var options = Options;
        var now = _time.GetUtcNow();

        if (success)
        {
            entry.RecordSuccess(now, latency);
        }
        else
        {
            entry.RecordFailure(
                now,
                options.BaseQuarantine,
                options.MaxQuarantine,
                options.FailuresBeforeQuarantine,
                error
            );
        }

        lock (_gate)
        {
            _strategy.Report(entry, success);
        }

        RaiseChanged();
    }

    /// <summary>Whether an endpoint passes the protocol and country filters.</summary>
    public static bool Matches(ProxyEndpoint endpoint, ProxyOptions options)
    {
        ArgumentNullException.ThrowIfNull(endpoint);
        ArgumentNullException.ThrowIfNull(options);

        if (
            options.Protocols != ProxyProtocolFilter.None
            && !options.Protocols.HasFlag(ProxyOptions.ToFlag(endpoint.Protocol))
        )
        {
            return false;
        }

        return options.Countries.Count == 0
            || (
                endpoint.Country is { } country && options.Countries.Contains(country, StringComparer.OrdinalIgnoreCase)
            );
    }

    private ProxyEntry? SelectAvailable()
    {
        lock (_gate)
        {
            var now = _time.GetUtcNow();

            // Best-first ordering: sticky selection takes the head of this list, and the other
            // strategies benefit from healthy proxies being clustered at the front.
            var candidates = _entries
                .Where(entry => entry.IsAvailable(now))
                .OrderByDescending(entry => entry.Health == ProxyHealthState.Alive)
                .ThenBy(entry => entry.Latency ?? TimeSpan.MaxValue)
                .ThenByDescending(entry => entry.Endpoint.Score)
                .ToArray();

            return _strategy.Pick(candidates);
        }
    }

    private void RaiseChanged() => Changed?.Invoke(this, EventArgs.Empty);
}