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

namespace PrivaPub.Infrastructure.Push
{
	// Web Push message encryption, on System.Security.Cryptography alone (completion program R9). Two encodings, as Mastodon
	// sends them:
	// - RFC 8291 (`aes128gcm`, in RFC 8188's record), for a subscription made with `standard` (Mastodon 4.4);
	// - the draft before it (`aesgcm`: draft-ietf-webpush-encryption-04 over draft-ietf-httpbis-encryption-encoding-03), for
	//   every other one, with its salt in the Encryption header and our public key in Crypto-Key.
	// Both: a fresh P-256 key pair and a fresh 16-byte salt per message, ECDH with the client's p256dh, HKDF-SHA-256 keyed by
	// its auth secret, AES-128-GCM, one record. Only the client can read the message; the push service carries it blind.
	public static class WebPushCrypto
	{
		public const int RecordSize = 4096;
		// the largest payload a 4096-byte aes128gcm body holds: 86 bytes of header, the delimiter and the tag
		public const int MaxPayload = RecordSize - 86 - 1 - 16;
		const int SaltLength = 16;
		const int TagLength = 16;

		public sealed record Message(byte[] Body, string ContentEncoding, byte[] Salt, byte[] ServerPublicKey);

		public static Message Aes128Gcm(ReadOnlySpan<byte> plaintext, byte[] clientPublicKey, byte[] authSecret)
		{
			using var server = ECDiffieHellman.Create(ECCurve.NamedCurves.nistP256);
			return Aes128Gcm(plaintext, clientPublicKey, authSecret, server, RandomNumberGenerator.GetBytes(SaltLength));
		}

		public static Message Aes128Gcm(ReadOnlySpan<byte> plaintext, byte[] clientPublicKey, byte[] authSecret, ECDiffieHellman server, byte[] salt)
		{
			if (plaintext.Length > MaxPayload)
				throw new ArgumentException($"a push payload holds {MaxPayload} bytes at most", nameof(plaintext));
			var serverPublic = PublicPoint(server);
			var shared = SharedSecret(server, clientPublicKey);
			// RFC 8291 3.3 and 3.4: PRK_key = HMAC-SHA-256(auth, ecdh); IKM = HKDF-Expand(PRK_key, "WebPush: info" 0x00
			// ua_public as_public, 32)
			var keyInfo = Concat(Encoding.ASCII.GetBytes("WebPush: info\0"), clientPublicKey, serverPublic);
			var ikm = HKDF.Expand(HashAlgorithmName.SHA256, HKDF.Extract(HashAlgorithmName.SHA256, shared, authSecret), 32, keyInfo);
			// RFC 8188 2.2 and 2.3: the content encryption key and the nonce, from the salt
			var prk = HKDF.Extract(HashAlgorithmName.SHA256, ikm, salt);
			var cek = HKDF.Expand(HashAlgorithmName.SHA256, prk, 16, Encoding.ASCII.GetBytes("Content-Encoding: aes128gcm\0"));
			var nonce = HKDF.Expand(HashAlgorithmName.SHA256, prk, 12, Encoding.ASCII.GetBytes("Content-Encoding: nonce\0"));
			// one record, the last: the payload, its delimiter 0x02, no padding
			var record = new byte[plaintext.Length + 1];
			plaintext.CopyTo(record);
			record[^1] = 0x02;
			var encrypted = Seal(cek, nonce, record);
			// RFC 8188 2.1's header: the salt, the record size, the key id's length and the key id (RFC 8291: our public key)
			var body = new byte[SaltLength + 4 + 1 + serverPublic.Length + encrypted.Length];
			salt.CopyTo(body, 0);
			BinaryPrimitives.WriteUInt32BigEndian(body.AsSpan(SaltLength), RecordSize);
			body[SaltLength + 4] = (byte)serverPublic.Length;
			serverPublic.CopyTo(body, SaltLength + 5);
			encrypted.CopyTo(body, SaltLength + 5 + serverPublic.Length);
			return new Message(body, "aes128gcm", salt, serverPublic);
		}

		public static Message AesGcm(ReadOnlySpan<byte> plaintext, byte[] clientPublicKey, byte[] authSecret)
		{
			using var server = ECDiffieHellman.Create(ECCurve.NamedCurves.nistP256);
			return AesGcm(plaintext, clientPublicKey, authSecret, server, RandomNumberGenerator.GetBytes(SaltLength));
		}

		public static Message AesGcm(ReadOnlySpan<byte> plaintext, byte[] clientPublicKey, byte[] authSecret, ECDiffieHellman server, byte[] salt)
		{
			if (plaintext.Length > MaxPayload)
				throw new ArgumentException($"a push payload holds {MaxPayload} bytes at most", nameof(plaintext));
			var serverPublic = PublicPoint(server);
			var shared = SharedSecret(server, clientPublicKey);
			// the draft's: IKM = HKDF(salt: auth, ecdh, "Content-Encoding: auth" 0x00, 32), then the key and nonce from the salt
			// over a context of both public keys, each after its length
			var ikm = HKDF.Expand(HashAlgorithmName.SHA256, HKDF.Extract(HashAlgorithmName.SHA256, shared, authSecret), 32,
				Encoding.ASCII.GetBytes("Content-Encoding: auth\0"));
			var context = Concat(Encoding.ASCII.GetBytes("P-256\0"), Length16(clientPublicKey), clientPublicKey, Length16(serverPublic), serverPublic);
			var prk = HKDF.Extract(HashAlgorithmName.SHA256, ikm, salt);
			var cek = HKDF.Expand(HashAlgorithmName.SHA256, prk, 16, Concat(Encoding.ASCII.GetBytes("Content-Encoding: aesgcm\0"), context));
			var nonce = HKDF.Expand(HashAlgorithmName.SHA256, prk, 12, Concat(Encoding.ASCII.GetBytes("Content-Encoding: nonce\0"), context));
			// two bytes saying how much padding follows (none), then the payload
			var padded = new byte[2 + plaintext.Length];
			plaintext.CopyTo(padded.AsSpan(2));
			return new Message(Seal(cek, nonce, padded), "aesgcm", salt, serverPublic);
		}

		/// <summary>A client's public key: an uncompressed point on P-256, or a CryptographicException.</summary>
		public static ECDiffieHellman ImportPublic(byte[] point)
		{
			if (point is not { Length: 65 } || point[0] != 0x04)
				throw new CryptographicException("not an uncompressed P-256 point");
			var key = ECDiffieHellman.Create(new ECParameters
			{
				Curve = ECCurve.NamedCurves.nistP256,
				Q = new ECPoint { X = point[1..33], Y = point[33..] }
			});
			return key;
		}

		/// <summary>Whether a subscription's keys can be encrypted to: a point on P-256 and a 16-byte secret.</summary>
		public static bool AreUsable(byte[] p256dh, byte[] auth)
		{
			if (auth is not { Length: 16 })
				return false;
			try
			{
				using var key = ImportPublic(p256dh);
				using var probe = ECDiffieHellman.Create(ECCurve.NamedCurves.nistP256);
				probe.DeriveRawSecretAgreement(key.PublicKey);
				return true;
			}
			catch (CryptographicException)
			{
				return false;
			}
		}

		public static byte[] PublicPoint(ECAlgorithm key)
		{
			var parameters = key.ExportParameters(false);
			return [0x04, .. parameters.Q.X, .. parameters.Q.Y];
		}

		public static string Base64Url(ReadOnlySpan<byte> bytes) =>
			Convert.ToBase64String(bytes).TrimEnd('=').Replace('+', '-').Replace('/', '_');

		/// <summary>base64url or base64, padded or not, as clients send keys; null when it is neither.</summary>
		public static byte[] FromBase64Url(string text)
		{
			if (string.IsNullOrWhiteSpace(text))
				return default;
			var normal = text.Trim().TrimEnd('=').Replace('-', '+').Replace('_', '/');
			normal = normal.PadRight(normal.Length + (4 - normal.Length % 4) % 4, '=');
			var bytes = new byte[normal.Length * 3 / 4];
			return Convert.TryFromBase64String(normal, bytes, out var written) ? bytes[..written] : default;
		}

		static byte[] SharedSecret(ECDiffieHellman server, byte[] clientPublicKey)
		{
			using var client = ImportPublic(clientPublicKey);
			return server.DeriveRawSecretAgreement(client.PublicKey);
		}

		static byte[] Seal(byte[] key, byte[] nonce, byte[] plaintext)
		{
			var output = new byte[plaintext.Length + TagLength];
			using var aes = new AesGcm(key, TagLength);
			aes.Encrypt(nonce, plaintext, output.AsSpan(0, plaintext.Length), output.AsSpan(plaintext.Length));
			return output;
		}

		static byte[] Length16(byte[] value)
		{
			var length = new byte[2];
			BinaryPrimitives.WriteUInt16BigEndian(length, (ushort)value.Length);
			return length;
		}

		static byte[] Concat(params byte[][] parts)
		{
			var all = new byte[parts.Sum(p => p.Length)];
			var at = 0;
			foreach (var part in parts)
			{
				part.CopyTo(all, at);
				at += part.Length;
			}
			return all;
		}
	}
}