using System.Buffers.Binary;
using System.Security.Cryptography;
using System.Text;

namespace Fleet.Identity;

/// <summary>
/// 节点这一侧用来确认「对面确实是我的操作员授权过的 Backend」的校验器：<c>Required</c> 为真（当前操作员集合里登记了服务器密钥）
/// 时，服务器的挑战必须带着对得上的签名，否则节点不继续认证、也不交出任何东西。
/// </summary>
public interface IServerVerifier
{
    bool Required { get; }

    /// <summary>服务器用它的服务器密钥对「挑战随机数加它这次连接的临时公钥」的签名。</summary>
    bool Verify(string nonce, ReadOnlySpan<byte> serverEphemeralSpki, ReadOnlySpan<byte> signature);
}

/// <summary>服务器在挑战里给出的身份证明：服务器密钥（操作员签进集合）对「随机数加临时公钥」签名。纯函数。</summary>
public static class ServerHello
{
    /// <summary>被签名的内容；服务器那一侧只拿到这个、用自己的密钥签（服务器密钥不一定是 BCL 的 ECDsa，可以在别处）。</summary>
    public static byte[] Payload(string nonce, ReadOnlySpan<byte> serverEphemeralSpki) =>
        Encoding.UTF8.GetBytes("fleet|v1|server|" + nonce + "|" + Convert.ToBase64String(serverEphemeralSpki));

    public static byte[] Sign(ECDsa serverKey, string nonce, ReadOnlySpan<byte> serverEphemeralSpki) =>
        serverKey.SignData(Payload(nonce, serverEphemeralSpki), HashAlgorithmName.SHA256, DSASignatureFormat.IeeeP1363FixedFieldConcatenation);

    public static bool Verify(ReadOnlySpan<byte> serverKeySpki, string nonce, ReadOnlySpan<byte> serverEphemeralSpki, ReadOnlySpan<byte> signature)
    {
        try
        {
            using ECDsa key = ECDsa.Create();
            key.ImportSubjectPublicKeyInfo(serverKeySpki, out _);
            return key.VerifyData(Payload(nonce, serverEphemeralSpki), signature, HashAlgorithmName.SHA256, DSASignatureFormat.IeeeP1363FixedFieldConcatenation);
        }
        catch (CryptographicException)
        {
            return false;
        }
    }
}

/// <summary>
/// 节点和 Backend 之间控制链路的应用层加密：两边的临时 ECDH 密钥协商出会话密钥，每个方向一把，AES-256-GCM，每帧带计数器。
/// 所以即使 TLS 在中间被终止（Cloudflare、被冒充的证书），中间人没有服务器密钥，就拿不到明文，也改不了任何一帧。
/// 帧格式：<c>计数器(8，大端) ‖ 密文 ‖ 标签(16)</c>；明文是 <c>帧类型(1) ‖ 内容</c>；计数器必须严格递增，重放、乱序、丢帧、篡改都会解密失败。
/// </summary>
public sealed class LinkCipher : IDisposable
{
    public const int Overhead = 8 + 16;

    private const byte NodeToServer = 0;
    private const byte ServerToNode = 1;

    private readonly AesGcm _send;
    private readonly AesGcm _receive;
    private readonly byte _sendDirection;
    private readonly byte _receiveDirection;
    private readonly object _gate = new();
    private long _sendCounter;
    private long _receiveCounter;

    private LinkCipher(byte[] sendKey, byte[] receiveKey, byte sendDirection, byte receiveDirection)
    {
        _send = new AesGcm(sendKey, 16);
        _receive = new AesGcm(receiveKey, 16);
        _sendDirection = sendDirection;
        _receiveDirection = receiveDirection;
    }

    public static LinkCipher ForNode(EphemeralKey own, ReadOnlySpan<byte> serverEphemeralSpki, string nonce)
    {
        (byte[] n2s, byte[] s2n) = Derive(own, serverEphemeralSpki, own.PublicKeySpki, serverEphemeralSpki, nonce);
        return new LinkCipher(n2s, s2n, NodeToServer, ServerToNode);
    }

    public static LinkCipher ForServer(EphemeralKey own, ReadOnlySpan<byte> nodeEphemeralSpki, string nonce)
    {
        (byte[] n2s, byte[] s2n) = Derive(own, nodeEphemeralSpki, nodeEphemeralSpki, own.PublicKeySpki, nonce);
        return new LinkCipher(s2n, n2s, ServerToNode, NodeToServer);
    }

    private static (byte[] NodeToServerKey, byte[] ServerToNodeKey) Derive(EphemeralKey own, ReadOnlySpan<byte> remote, ReadOnlySpan<byte> nodeEphemeral, ReadOnlySpan<byte> serverEphemeral, string nonce)
    {
        byte[] secret = own.AgreeRaw(remote);
        try
        {
            byte[] info = [.. Encoding.UTF8.GetBytes("fleet|v1|link|"), .. SHA256.HashData(serverEphemeral), .. SHA256.HashData(nodeEphemeral)];
            byte[] material = HKDF.DeriveKey(HashAlgorithmName.SHA256, secret, 64, Encoding.UTF8.GetBytes(nonce), info);
            return (material[..32], material[32..]);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(secret);
        }
    }

    /// <summary>加密一帧。</summary>
    public byte[] Seal(byte kind, ReadOnlySpan<byte> data)
    {
        long counter;
        lock (_gate)
        {
            counter = _sendCounter++;
        }

        byte[] output = new byte[Overhead + 1 + data.Length];
        BinaryPrimitives.WriteInt64BigEndian(output, counter);
        byte[] plaintext = new byte[1 + data.Length];
        plaintext[0] = kind;
        data.CopyTo(plaintext.AsSpan(1));
        _send.Encrypt(Nonce(counter), plaintext, output.AsSpan(8, plaintext.Length), output.AsSpan(8 + plaintext.Length, 16), Aad(_sendDirection, counter));
        return output;
    }

    /// <summary>解密一帧；计数器不是期望的下一个、或者标签不对，都返回 false。</summary>
    public bool TryOpen(ReadOnlySpan<byte> frame, out byte kind, out byte[] data)
    {
        kind = 0;
        data = [];
        if (frame.Length < Overhead + 1)
        {
            return false;
        }

        long counter = BinaryPrimitives.ReadInt64BigEndian(frame);
        lock (_gate)
        {
            if (counter != _receiveCounter)
            {
                return false;
            }

            int length = frame.Length - Overhead;
            byte[] plaintext = new byte[length];
            try
            {
                _receive.Decrypt(Nonce(counter), frame.Slice(8, length), frame.Slice(8 + length, 16), plaintext, Aad(_receiveDirection, counter));
            }
            catch (CryptographicException)
            {
                return false;
            }

            _receiveCounter++;
            kind = plaintext[0];
            data = plaintext[1..];
            return true;
        }
    }

    private static byte[] Nonce(long counter)
    {
        byte[] nonce = new byte[12];
        BinaryPrimitives.WriteInt64BigEndian(nonce.AsSpan(4), counter);
        return nonce;
    }

    private static byte[] Aad(byte direction, long counter)
    {
        byte[] aad = new byte[9];
        aad[0] = direction;
        BinaryPrimitives.WriteInt64BigEndian(aad.AsSpan(1), counter);
        return aad;
    }

    public void Dispose()
    {
        _send.Dispose();
        _receive.Dispose();
    }
}