using System.Collections.Concurrent;
using System.Net;
using System.Security.Cryptography;

namespace Fleet.Identity;

/// <summary>
/// 监听方的敲门关卡：新连接的前 32 字节必须是名册里某个成员的敲门包，否则不写任何字节、直接断开。
/// 每个来源地址限速（陌生人用随机数据也要消耗 HMAC 计算），nonce 在时间窗内去重（防重放）。
/// 名册由调用方提供（对端公钥列表）；配对密钥按公钥缓存，名册变了自然只算新的。
/// </summary>
public sealed class KnockGate
{
    private static readonly TimeSpan s_readTimeout = TimeSpan.FromSeconds(3);

    private readonly NodeIdentity _own;
    private readonly Func<IReadOnlyList<byte[]>> _peers;
    private readonly TimeProvider _time;
    private readonly int _attemptsPerPeriod;
    private readonly long _periodMs;
    private readonly ConcurrentDictionary<string, byte[]> _keys = new();
    private readonly object _gate = new();
    private readonly Dictionary<string, long> _seenNonces = [];
    private readonly Dictionary<string, (long WindowStartMs, int Count)> _sources = [];
    private long _lastSweepMs;

    public KnockGate(NodeIdentity own, Func<IReadOnlyList<byte[]>> peers, TimeProvider time, int attemptsPerPeriod = 300, TimeSpan? period = null)
    {
        _own = own;
        _peers = peers;
        _time = time;
        _attemptsPerPeriod = attemptsPerPeriod;
        _periodMs = (long)(period ?? TimeSpan.FromSeconds(10)).TotalMilliseconds;
    }

    /// <summary>读 32 字节并核对。通过返回 true，流从敲门包之后继续（TLS 或明文首帧）；不通过返回 false，调用方直接关闭连接，不要写任何东西。</summary>
    public async ValueTask<bool> AdmitAsync(Stream stream, IPAddress? source, CancellationToken ct)
    {
        long now = _time.GetUtcNow().ToUnixTimeMilliseconds();
        if (!TryCountSource(source, now))
        {
            return false;
        }

        byte[] packet = new byte[Knock.PacketBytes];
        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(ct);
        timeout.CancelAfter(s_readTimeout);
        try
        {
            await stream.ReadExactlyAsync(packet, timeout.Token);
        }
        catch (Exception ex) when (ex is EndOfStreamException or OperationCanceledException or IOException or ObjectDisposedException)
        {
            return false;
        }

        return Accept(packet, now);
    }

    /// <summary>核对一个敲门包（不读流）：名册里有成员的密钥对得上，并且 nonce 没用过。</summary>
    public bool Accept(ReadOnlySpan<byte> packet, long nowMs)
    {
        bool matched = false;
        foreach (byte[] peer in _peers())
        {
            byte[] key = _keys.GetOrAdd(Convert.ToBase64String(peer), _ => Knock.DeriveKey(_own, peer));
            matched |= Knock.Check(key, packet, nowMs);
        }

        if (!matched)
        {
            return false;
        }

        string nonce = Convert.ToHexString(packet[..Knock.NonceBytes]);
        lock (_gate)
        {
            Sweep(nowMs);
            return _seenNonces.TryAdd(nonce, nowMs + (Knock.WindowSeconds * 3 * 1000L));
        }
    }

    private bool TryCountSource(IPAddress? source, long nowMs)
    {
        string key = source is null ? "" : source.IsIPv4MappedToIPv6 ? source.MapToIPv4().ToString() : source.ToString();
        lock (_gate)
        {
            Sweep(nowMs);
            (long start, int count) = _sources.TryGetValue(key, out var entry) && nowMs - entry.WindowStartMs < _periodMs ? entry : (nowMs, 0);
            _sources[key] = (start, count + 1);
            return count < _attemptsPerPeriod;
        }
    }

    private void Sweep(long nowMs)
    {
        if (nowMs - _lastSweepMs < 5000 && _seenNonces.Count < 4096 && _sources.Count < 4096)
        {
            return;
        }

        _lastSweepMs = nowMs;
        foreach (string nonce in _seenNonces.Where(p => p.Value <= nowMs).Select(p => p.Key).ToList())
        {
            _seenNonces.Remove(nonce);
        }

        foreach (string source in _sources.Where(p => nowMs - p.Value.WindowStartMs >= _periodMs).Select(p => p.Key).ToList())
        {
            _sources.Remove(source);
        }
    }
}