using System.Globalization;
using System.Security.Cryptography;
using System.Text;
using SpawnDev.BlazorJS.Cryptography;
using static SpawnDev.BlazorJS.Cryptography.PortableCrypto;

namespace DataTransfer.Client.Services;

/// <summary>
///     End-to-end encryption of one connection. Runs in the browser (WebAssembly + WebCrypto); keys never leave it.
///     <list type="bullet">
///         <item>
///             Key agreement: ephemeral ECDH P-256; the private key is non-extractable and discarded after the
///             agreement.
///         </item>
///         <item>
///             Commit-then-reveal: the joiner first sends a hash commitment to its public key, the creator then reveals
///             its key, and only afterwards the joiner reveals its key, which the creator checks against the commitment.
///             A relay in the middle therefore has to fix the keys it shows each side before it learns the real keys, and
///             can't search for keys that produce matching safety codes.
///         </item>
///         <item>
///             One AES-256-GCM key per direction (PBKDF2-SHA-256 over the shared secret, salted with direction, code and
///             both public keys), so a frame can't be reflected back to its sender.
///         </item>
///         <item>
///             Every frame carries a sequence number inside the ciphertext: replayed, reordered or dropped frames are
///             rejected.
///         </item>
///         <item>
///             The safety code is a hash of both public keys; if both users see the same code, nobody intercepted the
///             agreement.
///         </item>
///     </list>
/// </summary>
public sealed class SecureChannel : IDisposable
{
    private const string Protocol = "DataTransfer/v3";
    private const int KeyDerivationIterations = 10_000;

    private readonly PortableCrypto _crypto;
    private PortableECDHKey? _keyPair;
    private PortableAESGCMKey? _sendKey;
    private PortableAESGCMKey? _receiveKey;
    private long _sendSequence;
    private long _receiveSequence;

    private SecureChannel(PortableCrypto crypto, PortableECDHKey keyPair, byte[] publicKey, byte[] commitment)
    {
        _crypto = crypto;
        _keyPair = keyPair;
        PublicKey = publicKey;
        Commitment = commitment;
    }

    /// <summary>Ephemeral ECDH public key (SPKI) to send to the peer.</summary>
    public byte[] PublicKey { get; }

    /// <summary>Hash commitment to <see cref="PublicKey" />; the joiner sends it before revealing its key.</summary>
    public byte[] Commitment { get; }

    public bool IsEstablished => _sendKey is not null;

    public static async Task<SecureChannel> CreateAsync(PortableCrypto crypto)
    {
        ArgumentNullException.ThrowIfNull(crypto);

        var keyPair = await crypto.GenerateECDHKey(NamedCurve.P256, false);
        try
        {
            var publicKey = await crypto.ExportPublicKeySpki(keyPair);
            var commitment = await ComputeCommitmentAsync(crypto, publicKey);
            return new SecureChannel(crypto, keyPair, publicKey, commitment);
        }
        catch
        {
            keyPair.Dispose();
            throw;
        }
    }

    /// <summary>
    ///     Completes the key agreement as creator and returns the safety code both users can compare.
    /// </summary>
    /// <param name="joinerCommitment">The commitment the joiner sent before the creator revealed its key.</param>
    /// <exception cref="CryptographicException">The joiner's key does not match its commitment.</exception>
    public async Task<string> EstablishAsCreatorAsync(byte[] joinerPublicKey, byte[] joinerCommitment, string code)
    {
        ArgumentNullException.ThrowIfNull(joinerPublicKey);
        ArgumentNullException.ThrowIfNull(joinerCommitment);

        var expected = await ComputeCommitmentAsync(_crypto, joinerPublicKey);
        if (!CryptographicOperations.FixedTimeEquals(expected, joinerCommitment))
        {
            throw new CryptographicException("The peer's public key does not match its commitment.");
        }

        return await EstablishAsync(joinerPublicKey, true, code);
    }

    /// <summary>
    ///     Completes the key agreement as joiner and returns the safety code. Must only be called after
    ///     <see cref="Commitment" /> was sent, and <see cref="PublicKey" /> must only be revealed afterwards.
    /// </summary>
    public Task<string> EstablishAsJoinerAsync(byte[] creatorPublicKey, string code) =>
        EstablishAsync(creatorPublicKey, false, code);

    /// <summary>Encrypts the next outgoing frame. Callers must send sealed frames in the order they were sealed.</summary>
    public Task<byte[]> SealAsync(Frame frame)
    {
        ArgumentNullException.ThrowIfNull(frame);

        var key = _sendKey ?? throw new InvalidOperationException("The channel is not established.");
        return _crypto.Encrypt(key, FrameCodec.Encode(++_sendSequence, frame));
    }

    /// <summary>Decrypts and verifies the next incoming frame.</summary>
    /// <exception cref="CryptographicException">The frame was tampered with, replayed, reordered or dropped.</exception>
    public async Task<Frame> OpenAsync(byte[] ciphertext)
    {
        ArgumentNullException.ThrowIfNull(ciphertext);

        var key = _receiveKey ?? throw new InvalidOperationException("The channel is not established.");

        byte[]? plaintext;
        try
        {
            plaintext = await _crypto.Decrypt(key, ciphertext);
        }
        catch (Exception ex) when (ex is not OperationCanceledException)
        {
            throw new CryptographicException("The frame could not be decrypted.", ex);
        }

        if (plaintext is null)
        {
            throw new CryptographicException("The frame could not be decrypted.");
        }

        long sequence;
        Frame frame;
        try
        {
            (sequence, frame) = FrameCodec.Decode(plaintext);
        }
        catch (InvalidDataException ex)
        {
            throw new CryptographicException("The frame is malformed.", ex);
        }

        if (sequence != _receiveSequence + 1)
        {
            throw new CryptographicException("A frame was replayed, reordered or dropped.");
        }

        _receiveSequence = sequence;
        return frame;
    }

    public void Dispose()
    {
        _keyPair?.Dispose();
        _sendKey?.Dispose();
        _receiveKey?.Dispose();
        _keyPair = null;
        _sendKey = null;
        _receiveKey = null;
    }

    private static async Task<byte[]> ComputeCommitmentAsync(PortableCrypto crypto, byte[] publicKey) =>
        await crypto.Digest(HashName.SHA256, Concat(Encoding.UTF8.GetBytes($"{Protocol}|commitment|"), publicKey));

    private static string FormatSafetyCode(byte[] digest)
    {
        var groups = new string[4];
        for (var group = 0; group < groups.Length; group++)
        {
            ulong value = 0;
            for (var i = 0; i < 5; i++)
            {
                value = (value << 8) | digest[group * 5 + i];
            }

            groups[group] = (value % 100_000).ToString("D5", CultureInfo.InvariantCulture);
        }

        return string.Join(' ', groups);
    }

    private static byte[] Concat(params byte[][] parts)
    {
        var result = new byte[parts.Sum(part => part.Length)];
        var offset = 0;
        foreach (var part in parts)
        {
            part.CopyTo(result, offset);
            offset += part.Length;
        }

        return result;
    }

    private async Task<string> EstablishAsync(byte[] peerPublicKey, bool isCreator, string code)
    {
        ArgumentNullException.ThrowIfNull(peerPublicKey);
        ArgumentException.ThrowIfNullOrEmpty(code);

        var keyPair = _keyPair;
        if (keyPair is null || IsEstablished)
        {
            throw new InvalidOperationException("The channel can only be established once.");
        }

        if (peerPublicKey.AsSpan().SequenceEqual(PublicKey))
        {
            throw new CryptographicException("The peer presented our own public key.");
        }

        _keyPair = null;
        byte[]? sharedSecret = null;
        try
        {
            using var peerKey = await _crypto.ImportECDHKey(peerPublicKey, NamedCurve.P256);
            sharedSecret = await _crypto.DeriveBits(keyPair, peerKey, 256);

            var creatorKey = isCreator ? PublicKey : peerPublicKey;
            var joinerKey = isCreator ? peerPublicKey : PublicKey;
            var creatorToJoiner = await DeriveKeyAsync(sharedSecret, "creator->joiner", code, creatorKey, joinerKey);
            var joinerToCreator = await DeriveKeyAsync(sharedSecret, "joiner->creator", code, creatorKey, joinerKey);

            _sendKey = isCreator ? creatorToJoiner : joinerToCreator;
            _receiveKey = isCreator ? joinerToCreator : creatorToJoiner;

            var digest = await _crypto.Digest(HashName.SHA256, Concat(Encoding.UTF8.GetBytes($"{Protocol}|safety-code|"), creatorKey, joinerKey));
            return FormatSafetyCode(digest);
        }
        finally
        {
            if (sharedSecret is not null)
            {
                CryptographicOperations.ZeroMemory(sharedSecret);
            }

            keyPair.Dispose();
        }
    }

    private async Task<PortableAESGCMKey> DeriveKeyAsync(byte[] sharedSecret, string direction, string code, byte[] creatorKey, byte[] joinerKey)
    {
        var salt = await _crypto.Digest(
            HashName.SHA256,
            Concat(Encoding.UTF8.GetBytes($"{Protocol}|{direction}|{code}|"), creatorKey, joinerKey));

        return await _crypto.GenerateAESGCMKey(
            sharedSecret,
            salt,
            KeyDerivationIterations,
            HashName.SHA256,
            32,
            16,
            12,
            false);
    }
}