using System.Net;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Text.Json;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Rpc;

namespace Fleet.Link;

/// <summary>一个已经认证的点对点会话：对端是谁、数据面使用的会话凭据、对端地址。</summary>
public sealed record PeerSessionInfo(string PeerId, SessionCredentials Credentials, EndPoint? Remote, string? Version = null);

public sealed class PeerListenerOptions
{
    public required IPEndPoint Listen { get; init; }

    public required X509Certificate2 Certificate { get; init; }

    /// <summary>按对端 ID 查找允许连接的公钥（SPKI）。不在列表里的对端一律拒绝（期望状态里的 peers）。</summary>
    public required Func<string, byte[]?> LookupPeerKey { get; init; }

    public int MaxFrameBytes { get; init; } = 8 << 20;

    public int MaxSessions { get; init; } = 8;

    public TimeSpan HandshakeTimeout { get; init; } = TimeSpan.FromSeconds(10);
}

/// <summary>
/// 监听点对点会话（Agent 一侧）：TLS（自签名证书，对端按 SPKI 指纹固定）→ 挑战 → 对端用身份私钥签名（purpose = peer）
/// → 临时 ECDH 派生数据面凭据。认证通过之后把会话交给产品代码，由它挂上自己的路由。
/// </summary>
public sealed class PeerListener
{
    private readonly PeerListenerOptions _options;
    private readonly Func<PeerSessionInfo, IRpcTransport, CancellationToken, Task> _onSession;
    private readonly INodeLog _log;
    private readonly TimeProvider _time;
    private int _sessions;

    public PeerListener(PeerListenerOptions options, Func<PeerSessionInfo, IRpcTransport, CancellationToken, Task> onSession, INodeLog log, TimeProvider time)
    {
        _options = options;
        _onSession = onSession;
        _log = log;
        _time = time;
    }

    /// <summary>实际监听的端口（配置为 0 时由系统分配）；<see cref="RunAsync"/> 开始监听之后有效。</summary>
    public int Port { get; private set; }

    public TaskCompletionSource Started { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);

    public async Task RunAsync(CancellationToken ct)
    {
        var listener = new TcpListener(_options.Listen);
        listener.Start();
        Port = ((IPEndPoint)listener.LocalEndpoint).Port;
        Started.TrySetResult();
        try
        {
            while (!ct.IsCancellationRequested)
            {
                TcpClient client = await listener.AcceptTcpClientAsync(ct);
                if (Interlocked.Increment(ref _sessions) > _options.MaxSessions)
                {
                    Interlocked.Decrement(ref _sessions);
                    client.Dispose();
                    continue;
                }

                _ = Task.Run(() => HandleAsync(client, ct), CancellationToken.None);
            }
        }
        catch (OperationCanceledException)
        {
            // 停止监听
        }
        finally
        {
            listener.Stop();
        }
    }

    private async Task HandleAsync(TcpClient client, CancellationToken ct)
    {
        try
        {
            client.NoDelay = true;
            PeerSockets.EnableKeepAlive(client.Client);
            using var timeout = new CancellationTokenSource(_options.HandshakeTimeout, _time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
            var tls = new SslStream(client.GetStream(), leaveInnerStreamOpen: false);
            await tls.AuthenticateAsServerAsync(
                new SslServerAuthenticationOptions { ServerCertificate = _options.Certificate, ClientCertificateRequired = false, EnabledSslProtocols = System.Security.Authentication.SslProtocols.Tls12 | System.Security.Authentication.SslProtocols.Tls13 },
                linked.Token);
            var transport = new TlsFrameTransport(tls, _options.MaxFrameBytes);
            PeerSessionInfo? info = await HandshakeAsync(transport, client.Client.RemoteEndPoint, linked.Token);
            if (info is null)
            {
                await transport.DisposeAsync();
                return;
            }

            _log.Info("peer session", ("probe", info.PeerId), ("result", "established"));
            await using (transport)
            {
                await _onSession(info, transport, ct);
            }
        }
        catch (Exception ex) when (ex is IOException or System.Security.Authentication.AuthenticationException or OperationCanceledException or SocketException or ObjectDisposedException)
        {
            // 握手失败或连接断开：这个会话结束，监听继续
        }
        finally
        {
            client.Dispose();
            Interlocked.Decrement(ref _sessions);
        }
    }

    private async Task<PeerSessionInfo?> HandshakeAsync(TlsFrameTransport transport, EndPoint? remote, CancellationToken ct)
    {
        using var ephemeral = new EphemeralKey();
        byte[] nonceBytes = RandomNumberGenerator.GetBytes(32);
        string nonce = Convert.ToBase64String(nonceBytes);
        string serverPub = Convert.ToBase64String(ephemeral.PublicKeySpki);
        await SendAsync(transport, RpcMessage.Notification(PeerMethods.Challenge, new PeerChallengeParams { Nonce = nonce, EcdhPublicKey = serverPub }, PeerRpcContext.Default.PeerChallengeParams), ct);

        RpcFrame? frame = await transport.ReceiveAsync(ct);
        RpcMessage? message = frame is { Kind: FrameKind.Text } f ? JsonSerializer.Deserialize(f.Data.Span, FleetRpcContext.Default.RpcMessage) : null;
        if (message?.Method != PeerMethods.Auth || message.Id is null || message.Params is null)
        {
            return null;
        }

        PeerAuthParams? auth = null;
        try
        {
            auth = message.Params.Value.Deserialize(PeerRpcContext.Default.PeerAuthParams);
        }
        catch (JsonException)
        {
            // 按缺失处理
        }

        byte[]? key = auth is null ? null : _options.LookupPeerKey(auth.PeerId);
        byte[] clientPublic = [];
        bool ok = key is not null
            && TryDecode(auth!.Signature, out byte[] signature)
            && TryDecode(auth.EcdhPublicKey, out clientPublic)
            && NodeIdentity.Verify(key, signature, SignaturePurpose.Peer, auth.PeerId, nonce, auth.EcdhPublicKey, serverPub);
        if (!ok)
        {
            // 对端不在允许列表里和签名错误返回同样的错误
            await SendAsync(transport, RpcMessage.Failure(message.Id, new FleetError(FleetErrors.Unauthorized, "peer authentication failed")), ct);
            _log.Warn("peer rejected", ("result", "unauthorized"));
            return null;
        }

        SessionCredentials credentials;
        try
        {
            credentials = ephemeral.Derive(clientPublic, nonceBytes, clientPublic, ephemeral.PublicKeySpki);
        }
        catch (CryptographicException)
        {
            await SendAsync(transport, RpcMessage.Failure(message.Id, new FleetError(FleetErrors.InvalidParams, "invalid ecdh key")), ct);
            return null;
        }

        await SendAsync(transport, RpcMessage.Success(message.Id.Value, new PeerAuthResult { SessionId = Convert.ToHexStringLower(credentials.SessionId) }, PeerRpcContext.Default.PeerAuthResult), ct);
        return new PeerSessionInfo(auth!.PeerId, credentials, remote, auth.Version is { Length: <= 64 } version ? version : null);
    }

    private static ValueTask SendAsync(TlsFrameTransport transport, RpcMessage message, CancellationToken ct) =>
        transport.SendAsync(new RpcFrame(FrameKind.Text, JsonSerializer.SerializeToUtf8Bytes(message, FleetRpcContext.Default.RpcMessage)), ct);

    private static bool TryDecode(string text, out byte[] data)
    {
        try
        {
            data = Convert.FromBase64String(text);
            return data.Length > 0;
        }
        catch (FormatException)
        {
            data = [];
            return false;
        }
    }
}

/// <summary>发起点对点会话（Probe 一侧）：按 SPKI 指纹固定对端证书，用身份私钥签名，派生同样的数据面凭据。</summary>
/// <summary>点对点会话的套接字：开启 TCP keepalive（空闲 15 秒后每 5 秒探测，连续 3 次无响应就断开，约 30 秒），主机丢失或路径被切断时连接不会一直挂着。</summary>
internal static class PeerSockets
{
    public static void EnableKeepAlive(Socket socket)
    {
        try
        {
            socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);
            socket.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.TcpKeepAliveTime, 15);
            socket.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.TcpKeepAliveInterval, 5);
            socket.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.TcpKeepAliveRetryCount, 3);
        }
        catch (Exception ex) when (ex is SocketException or ArgumentException or PlatformNotSupportedException)
        {
            // 个别平台不支持其中某个选项：退回系统默认
        }
    }
}

public static class PeerConnector
{
    public static async Task<(PeerSessionInfo Info, TlsFrameTransport Transport)> ConnectAsync(
        string host, int port, byte[] expectedSpkiSha256, NodeIdentity identity, string peerId, string? version, int maxFrameBytes, TimeProvider time, CancellationToken ct)
    {
        var client = new TcpClient { NoDelay = true };
        PeerSockets.EnableKeepAlive(client.Client);
        try
        {
            using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(10), time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
            await client.ConnectAsync(host, port, linked.Token);
            var tls = new SslStream(
                client.GetStream(), leaveInnerStreamOpen: false,
                (_, certificate, _, _) => certificate is not null && SpkiMatches(certificate, expectedSpkiSha256));
            await tls.AuthenticateAsClientAsync(new SslClientAuthenticationOptions { TargetHost = host }, linked.Token);
            var transport = new TlsFrameTransport(tls, maxFrameBytes);

            RpcFrame? first = await transport.ReceiveAsync(linked.Token);
            RpcMessage? challengeMessage = first is { Kind: FrameKind.Text } f ? JsonSerializer.Deserialize(f.Data.Span, FleetRpcContext.Default.RpcMessage) : null;
            if (challengeMessage?.Method != PeerMethods.Challenge || challengeMessage.Params is null)
            {
                throw new InvalidOperationException("the agent did not send a challenge");
            }

            PeerChallengeParams challenge = challengeMessage.Params.Value.Deserialize(PeerRpcContext.Default.PeerChallengeParams)!;
            using var ephemeral = new EphemeralKey();
            string clientPub = Convert.ToBase64String(ephemeral.PublicKeySpki);
            string signature = Convert.ToBase64String(identity.Sign(SignaturePurpose.Peer, peerId, challenge.Nonce, clientPub, challenge.EcdhPublicKey));
            RpcMessage request = RpcMessage.Request(1, PeerMethods.Auth, new PeerAuthParams { PeerId = peerId, EcdhPublicKey = clientPub, Signature = signature, Version = version }, PeerRpcContext.Default.PeerAuthParams);
            await transport.SendAsync(new RpcFrame(FrameKind.Text, JsonSerializer.SerializeToUtf8Bytes(request, FleetRpcContext.Default.RpcMessage)), linked.Token);

            RpcFrame? replyFrame = await transport.ReceiveAsync(linked.Token);
            RpcMessage? reply = replyFrame is { Kind: FrameKind.Text } r ? JsonSerializer.Deserialize(r.Data.Span, FleetRpcContext.Default.RpcMessage) : null;
            if (reply is null || reply.Error is not null)
            {
                throw new InvalidOperationException("peer authentication failed: " + reply?.Error?.Data?.Code);
            }

            SessionCredentials credentials = ephemeral.Derive(Convert.FromBase64String(challenge.EcdhPublicKey), Convert.FromBase64String(challenge.Nonce), ephemeral.PublicKeySpki, Convert.FromBase64String(challenge.EcdhPublicKey));
            string announced = reply.Result!.Value.Deserialize(PeerRpcContext.Default.PeerAuthResult)!.SessionId;
            if (announced != Convert.ToHexStringLower(credentials.SessionId))
            {
                throw new InvalidOperationException("session id mismatch: the key agreement did not produce the same credentials");
            }

            return (new PeerSessionInfo(peerId, credentials, client.Client.RemoteEndPoint), transport);
        }
        catch
        {
            client.Dispose();
            throw;
        }
    }

    /// <summary>固定证书：只看 SPKI 的 SHA-256 指纹，不看证书链和主机名（自签名证书）。</summary>
    private static bool SpkiMatches(X509Certificate certificate, byte[] expected)
    {
        X509Certificate2 full = certificate as X509Certificate2 ?? X509CertificateLoader.LoadCertificate(certificate.GetRawCertData());
        return CryptographicOperations.FixedTimeEquals(NodeCertificate.SpkiFingerprint(full), expected);
    }
}