using System.Collections.Concurrent;
using System.Net;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Authentication;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using Fleet.Identity;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>拨号方的一条数据面 TCP 连接（已经完成 TLS 和首帧认证）。</summary>
public sealed class TcpLink : IDisposable
{
    private readonly TcpClient _client;

    private TcpLink(TcpClient client, Stream stream, long tcpUs, long? tlsUs)
    {
        _client = client;
        Socket = client.Client;
        Stream = stream;
        TcpUs = tcpUs;
        TlsUs = tlsUs;
    }

    public Stream Stream { get; }

    /// <summary>底层套接字（内核采样用）。</summary>
    public Socket Socket { get; }

    public long TcpUs { get; }

    public long? TlsUs { get; }

    /// <summary>
    /// 连接 → （TLS）→ 认证首帧。<paramref name="sendAuth"/> 为 false 时只测量建立连接的耗时（connect）。
    /// 对端证书按 <see cref="BlockEnv.PeerSpkiSha256"/> 固定；没有指纹（目标任务，目标是任意地址）就只测握手耗时、不校验证书。
    /// 失败时抛出 <see cref="SocketException"/>、<see cref="AuthenticationException"/>、<see cref="IOException"/> 或超时的 <see cref="OperationCanceledException"/>。
    /// </summary>
    public static async Task<TcpLink> OpenAsync(BlockEnv env, int port, bool tls, uint streamId, bool sendAuth, CancellationToken ct)
    {
        var client = new TcpClient(env.PeerAddress!.AddressFamily) { NoDelay = true };
        try
        {
            long start = env.Clock.NowUs();
            await client.ConnectAsync(env.PeerAddress!, port, ct);
            long tcpUs = env.Clock.NowUs() - start;
            Stream stream = client.GetStream();
            if (env.KnockKey is { } knockKey)
            {
                // 预握手敲门：监听方不认就不会回任何字节
                await stream.WriteAsync(Knock.Build(knockKey, env.Time.GetUtcNow().ToUnixTimeMilliseconds()), ct);
            }

            long? tlsUs = null;
            if (tls)
            {
                long tlsStart = env.Clock.NowUs();
                byte[]? pin = env.PeerSpkiSha256;
#pragma warning disable CA5359 // 没有指纹时只测量握手耗时（目标任务的目标是任意地址）
                var ssl = new SslStream(stream, false, (_, certificate, _, _) => pin is null || (certificate is not null && Pinned(certificate, pin)));
#pragma warning restore CA5359
                await ssl.AuthenticateAsClientAsync(new SslClientAuthenticationOptions { TargetHost = env.PeerAddress!.ToString() }, ct);
                stream = ssl;
                tlsUs = env.Clock.NowUs() - tlsStart;
            }

            if (sendAuth)
            {
                DataSession session = env.Session() ?? throw new InvalidOperationException("no data session");
                byte[] auth = new byte[TcpWire.AuthBytes];
                Span<byte> nonce = stackalloc byte[TcpWire.NonceBytes];
                RandomNumberGenerator.Fill(nonce);
                TcpWire.WriteAuth(auth, session.SessionId, streamId, nonce, session.Key);
                await stream.WriteAsync(auth, ct);
            }

            return new TcpLink(client, stream, tcpUs, tlsUs);
        }
        catch
        {
            client.Dispose();
            throw;
        }
    }

    public void Dispose()
    {
        Stream.Dispose();
        _client.Dispose();
    }

    private static bool Pinned(X509Certificate certificate, byte[] expected)
    {
        X509Certificate2 full = certificate as X509Certificate2 ?? X509CertificateLoader.LoadCertificate(certificate.GetRawCertData());
        return CryptographicOperations.FixedTimeEquals(SHA256.HashData(full.PublicKey.ExportSubjectPublicKeyInfo()), expected);
    }
}

/// <summary>监听方收到的一条已经认证的数据面 TCP 连接。</summary>
public sealed class TcpConnection
{
    public TcpConnection(Stream stream, Socket socket, DataSession session, uint streamId)
    {
        Stream = stream;
        Socket = socket;
        Session = session;
        StreamId = streamId;
    }

    public Stream Stream { get; }

    /// <summary>底层套接字（内核采样用）。</summary>
    public Socket Socket { get; }

    public DataSession Session { get; }

    public uint StreamId { get; }
}

public delegate Task TcpResponder(TcpConnection connection, CancellationToken ct);

/// <summary>
/// 监听方的 TCP 数据面：每个端口一个监听器。连接先过预握手敲门，然后第一个字节是 0x16 就先做 TLS（节点自己的证书），否则直接是明文。
/// 随后是认证首帧（<c>"MN" | ver | session_id | stream_id | nonce | mac</c>，随机数只能用一次），通过之后交给按 <c>stream_id</c> 登记的处理函数。
/// 首帧在限定时间内没有到达或者校验失败，直接关闭连接并计数。
/// </summary>
public sealed class TcpAcceptor
{
    private static readonly TimeSpan s_authTimeout = TimeSpan.FromSeconds(5);

    private readonly IPAddress _bind;
    private readonly IReadOnlyList<int> _ports;
    private readonly X509Certificate2 _certificate;
    private readonly IDataSessions _sessions;
    private readonly DataPlaneCounters _counters;
    private readonly KnockGate _knock;
    private readonly ConcurrentDictionary<uint, TcpResponder> _handlers = new();
    private readonly List<TcpListener> _listeners = [];

    public TcpAcceptor(IPAddress bind, IReadOnlyList<int> ports, X509Certificate2 certificate, IDataSessions sessions, DataPlaneCounters counters, KnockGate knock)
    {
        _knock = knock;
        _bind = bind;
        _ports = ports;
        _certificate = certificate;
        _sessions = sessions;
        _counters = counters;
    }

    public IReadOnlyList<int> BoundPorts { get; private set; } = [];

    public void Register(uint streamId, TcpResponder handler) => _handlers[streamId] = handler;

    public void Unregister(uint streamId) => _handlers.TryRemove(streamId, out _);

    /// <summary>开始监听（端口被占用等错误在这里抛出）。</summary>
    public void Bind()
    {
        var bound = new List<int>();
        foreach (int port in _ports)
        {
            var listener = new TcpListener(_bind, port);
            if (_bind.Equals(IPAddress.IPv6Any))
            {
                listener.Server.DualMode = true;
            }

            listener.Start(128);
            _listeners.Add(listener);
            bound.Add(((IPEndPoint)listener.LocalEndpoint).Port);
        }

        BoundPorts = bound;
    }

    public async Task RunAsync(CancellationToken ct)
    {
        try
        {
            await Task.WhenAll(_listeners.Select(l => AcceptLoopAsync(l, ct)));
        }
        finally
        {
            foreach (TcpListener listener in _listeners)
            {
                listener.Stop();
            }
        }
    }

    private async Task AcceptLoopAsync(TcpListener listener, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            TcpClient client;
            try
            {
                client = await listener.AcceptTcpClientAsync(ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
            catch (SocketException)
            {
                continue;
            }

            _ = Task.Run(() => HandleAsync(client, ct), CancellationToken.None);
        }
    }

    private async Task HandleAsync(TcpClient client, CancellationToken ct)
    {
        using (client)
        {
            client.NoDelay = true;
            try
            {
                Stream stream = client.GetStream();
                using var authCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                authCts.CancelAfter(s_authTimeout);

                // 预握手敲门：名册里的 拨号方 才算得出敲门包；不通过就不写任何字节、直接断开
                if (!await _knock.AdmitAsync(stream, (client.Client.RemoteEndPoint as IPEndPoint)?.Address, authCts.Token))
                {
                    _counters.DropInvalid();
                    return;
                }

                byte[] first = new byte[1];
                if (await stream.ReadAsync(first, authCts.Token) != 1)
                {
                    return;
                }

                stream = new PrefixedStream(stream, first[0]);
                if (first[0] == 0x16)
                {
                    var ssl = new SslStream(stream, leaveInnerStreamOpen: false);
                    await ssl.AuthenticateAsServerAsync(new SslServerAuthenticationOptions { ServerCertificate = _certificate }, authCts.Token);
                    stream = ssl;
                }

                byte[] auth = new byte[TcpWire.AuthBytes];
                await stream.ReadExactlyAsync(auth, authCts.Token);
                if (!TcpWire.TryParseAuth(auth, out long sessionId, out uint streamId, out ReadOnlySpan<byte> nonce)
                    || _sessions.Find(sessionId) is not { } session
                    || !_handlers.TryGetValue(streamId, out TcpResponder? handler)
                    || !TcpWire.VerifyAuth(auth, session.Key)
                    || !session.UseNonce(nonce))
                {
                    _counters.DropInvalid();
                    return;
                }

                await handler(new TcpConnection(stream, client.Client, session, streamId), ct);
            }
            catch (Exception ex) when (ex is IOException or SocketException or OperationCanceledException or ObjectDisposedException or AuthenticationException or InvalidDataException or EndOfStreamException)
            {
                // 连接被对端关闭、超时或者格式错误：直接结束这个连接
            }
        }
    }

    /// <summary>把已经读出的第一个字节放回流的开头（用于探测是 TLS 还是明文）。</summary>
    private sealed class PrefixedStream : Stream
    {
        private readonly Stream _inner;
        private byte? _prefix;

        public PrefixedStream(Stream inner, byte prefix)
        {
            _inner = inner;
            _prefix = prefix;
        }

        public override bool CanRead => true;

        public override bool CanSeek => false;

        public override bool CanWrite => true;

        public override long Length => throw new NotSupportedException();

        public override long Position
        {
            get => throw new NotSupportedException();
            set => throw new NotSupportedException();
        }

        public override void Flush() => _inner.Flush();

        public override Task FlushAsync(CancellationToken cancellationToken) => _inner.FlushAsync(cancellationToken);

        public override int Read(byte[] buffer, int offset, int count) => Read(buffer.AsSpan(offset, count));

        public override int Read(Span<byte> buffer)
        {
            if (buffer.IsEmpty)
            {
                return 0;
            }

            if (_prefix is { } prefix)
            {
                buffer[0] = prefix;
                _prefix = null;
                return 1;
            }

            return _inner.Read(buffer);
        }

        public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
        {
            if (buffer.IsEmpty)
            {
                return ValueTask.FromResult(0);
            }

            if (_prefix is { } prefix)
            {
                buffer.Span[0] = prefix;
                _prefix = null;
                return ValueTask.FromResult(1);
            }

            return _inner.ReadAsync(buffer, cancellationToken);
        }

        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => ReadAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();

        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();

        public override void SetLength(long value) => throw new NotSupportedException();

        public override void Write(byte[] buffer, int offset, int count) => _inner.Write(buffer, offset, count);

        public override void Write(ReadOnlySpan<byte> buffer) => _inner.Write(buffer);

        public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default) => _inner.WriteAsync(buffer, cancellationToken);

        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => _inner.WriteAsync(buffer, offset, count, cancellationToken);

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                _inner.Dispose();
            }

            base.Dispose(disposing);
        }
    }
}