using System.Security.Cryptography;

namespace IdentityShroud.Core.Security;

public static class Encryption
{
    private readonly record struct AlgVersion(int Version, int NonceSize, int TagSize);

    private static AlgVersion[] _versions = 
    [
        new(0, 0, 0), // version 0 does not realy exist
        new(1, 12, 16), // version 1
    ];
    
    public static byte[] Encrypt(ReadOnlySpan<byte> plaintext, ReadOnlySpan<byte> key)
    {
        const int versionNumber = 1; 
        AlgVersion versionParams = _versions[versionNumber];

        int resultSize = 1 + versionParams.NonceSize + versionParams.TagSize + plaintext.Length;
        // allocate buffer for complete response 
        var result = new byte[resultSize];
        
        result[0] = (byte)versionParams.Version;

        // make the spans that point to the parts of the result where their data is located
        var nonce = result.AsSpan(1, versionParams.NonceSize); 
        var tag = result.AsSpan(1 + versionParams.NonceSize, versionParams.TagSize);
        var cipher = result.AsSpan(1 + versionParams.NonceSize + versionParams.TagSize);

        // use the spans to place the data directly in its place
        RandomNumberGenerator.Fill(nonce);
        using var aes = new AesGcm(key, versionParams.TagSize);
        aes.Encrypt(nonce, plaintext, cipher, tag);
        return result;
    }
    
    public static void Decrypt(ReadOnlyMemory<byte> input, ReadOnlySpan<byte> key, Span<byte> output)
    {
        AlgVersion versionParams = GetVersionParams(input);
        if (input.Length < 1 + versionParams.NonceSize + versionParams.TagSize)
            throw new ArgumentException("Cypher data is too short to be valid.", nameof(input));

        var payload = input.Span;
        ReadOnlySpan<byte> nonce = payload.Slice(1, versionParams.NonceSize);
        ReadOnlySpan<byte> tag = payload.Slice(1 + versionParams.NonceSize, versionParams.TagSize);
        ReadOnlySpan<byte> cipher = payload.Slice(1 + versionParams.NonceSize + versionParams.TagSize);

        using var aes = new AesGcm(key, versionParams.TagSize);
        try
        {
            aes.Decrypt(nonce, cipher, tag, output);
        }
        catch (CryptographicException ex)
        {
            // Tag verification failed → tampering or wrong key/nonce.
            throw new InvalidOperationException("Decryption failed – authentication tag mismatch.", ex);
        }
    }

    public static int GetDecryptedLength(ReadOnlyMemory<byte> input)
    {
        AlgVersion versionParams = GetVersionParams(input);
        int length = input.Length - (1 + versionParams.NonceSize + versionParams.TagSize);
        if (length < 0)
            throw new ArgumentException("Cypher data is too short to be valid.", nameof(input));
        
        return length;
    }

    private static AlgVersion GetVersionParams(ReadOnlyMemory<byte> input)
    {
        var versionNumber = (int)input.Span[0];
        if (versionNumber != 1)
            throw new ArgumentException("Invalid payload");
        
        return _versions[versionNumber];
    }
}