package com.dangerousthings.aliro;

import javacard.security.AESKey;
import javacard.security.KeyBuilder;
import javacardx.crypto.AEADCipher;
import javacardx.crypto.Cipher;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;

/**
 * Tests for {@link AliroGcm} — our userland AES-256-GCM, built on the JC
 * single-block AES-ECB primitive. Needed because the target card (NXP J3R180)
 * does not expose {@code AEADCipher.ALG_AES_GCM}. The NIST test vectors come
 * from {@code gcmEncryptExtIV256.rsp} (the CMVP AEAD AES-256 suite).
 */
class AliroGcmTest {

    @Test
    void nistTestCase13_emptyMessage_emptyAad_returnsTagOnly() {
        // K=0^256, IV=0^96, P=A=empty, expected T = 530f8afbc74536b9a963b4f1c4cb738b
        byte[] key = new byte[32];
        byte[] iv = new byte[12];
        byte[] pt = new byte[0];
        byte[] expectedTag = hex("530f8afbc74536b9a963b4f1c4cb738b");

        byte[] out = new byte[16];
        AliroGcm gcm = new AliroGcm();
        short n = gcm.encrypt(key, (short) 0, iv, (short) 0, pt, (short) 0, (short) 0, out, (short) 0);

        assertEquals(16, n, "empty plaintext → 16 bytes output (just the tag)");
        assertArrayEquals(expectedTag, out, "NIST GCM AES-256 Test Case 13 tag");
    }

    @Test
    void nistTestCase14_singleZeroBlock_returnsCiphertextAndTag() {
        // K=0^256, IV=0^96, P=0^128, A=empty
        // C = cea7403d4d606b6e074ec5d3baf39d18
        // T = d0d1c8a799996bf0265b98b5d48ab919
        byte[] key = new byte[32];
        byte[] iv = new byte[12];
        byte[] pt = new byte[16];
        byte[] expectedCt = hex("cea7403d4d606b6e074ec5d3baf39d18");
        byte[] expectedTag = hex("d0d1c8a799996bf0265b98b5d48ab919");

        byte[] out = new byte[32];
        AliroGcm gcm = new AliroGcm();
        short n = gcm.encrypt(key, (short) 0, iv, (short) 0, pt, (short) 0, (short) pt.length, out, (short) 0);

        assertEquals(32, n, "16B plaintext + 16B tag");

        byte[] ct = new byte[16];
        System.arraycopy(out, 0, ct, 0, 16);
        byte[] tag = new byte[16];
        System.arraycopy(out, 16, tag, 0, 16);

        assertArrayEquals(expectedCt, ct, "NIST GCM AES-256 Test Case 14 ciphertext");
        assertArrayEquals(expectedTag, tag, "NIST GCM AES-256 Test Case 14 tag");
    }

    @Test
    void nistTestCase15_fourBlocks_realKey() {
        // K  = feffe9928665731c6d6a8f9467308308 feffe9928665731c6d6a8f9467308308
        // IV = cafebabefacedbaddecaf888
        // P  = 4 blocks, A = empty
        byte[] key = hex("feffe9928665731c6d6a8f9467308308"
                       + "feffe9928665731c6d6a8f9467308308");
        byte[] iv  = hex("cafebabefacedbaddecaf888");
        byte[] pt  = hex("d9313225f88406e5a55909c5aff5269a"
                       + "86a7a9531534f7da2e4c303d8a318a72"
                       + "1c3c0c95956809532fcf0e2449a6b525"
                       + "b16aedf5aa0de657ba637b391aafd255");
        byte[] expectedCt  = hex("522dc1f099567d07f47f37a32a84427d"
                               + "643a8cdcbfe5c0c97598a2bd2555d1aa"
                               + "8cb08e48590dbb3da7b08b1056828838"
                               + "c5f61e6393ba7a0abcc9f662898015ad");
        byte[] expectedTag = hex("b094dac5d93471bdec1a502270e3cc6c");

        byte[] out = new byte[pt.length + 16];
        AliroGcm gcm = new AliroGcm();
        short n = gcm.encrypt(key, (short) 0, iv, (short) 0, pt, (short) 0, (short) pt.length, out, (short) 0);

        assertEquals(pt.length + 16, n);

        byte[] ct = new byte[pt.length];
        System.arraycopy(out, 0, ct, 0, pt.length);
        byte[] tag = new byte[16];
        System.arraycopy(out, pt.length, tag, 0, 16);

        assertArrayEquals(expectedCt, ct, "NIST GCM AES-256 Test Case 15 ciphertext");
        assertArrayEquals(expectedTag, tag, "NIST GCM AES-256 Test Case 15 tag");
    }

    /**
     * End-to-end: encrypt with AliroGcm, decrypt with jcardsim's AEADCipher
     * (available in tests via the patched jcardsim master in m2 volume).
     * Confirms wire-compatibility with the standard implementation.
     */
    @Test
    void roundTripAgainstJcardsimAEADCipher() {
        byte[] key = new byte[32];
        for (int i = 0; i < 32; i++) key[i] = (byte) (i * 7 + 3);
        byte[] iv = new byte[12];
        for (int i = 0; i < 12; i++) iv[i] = (byte) (0xA0 ^ i);
        byte[] pt = new byte[137];  // realistic Aliro Table 8-11 size
        for (int i = 0; i < pt.length; i++) pt[i] = (byte) (i ^ 0x5C);

        byte[] out = new byte[pt.length + 16];
        AliroGcm gcm = new AliroGcm();
        gcm.encrypt(key, (short) 0, iv, (short) 0,
            pt, (short) 0, (short) pt.length, out, (short) 0);

        // Decrypt with jcardsim's AEADCipher.
        AESKey k = (AESKey) KeyBuilder.buildKey(
            KeyBuilder.TYPE_AES, KeyBuilder.LENGTH_AES_256, false);
        k.setKey(key, (short) 0);
        AEADCipher jcGcm = (AEADCipher) Cipher.getInstance(AEADCipher.ALG_AES_GCM, false);
        jcGcm.init(k, Cipher.MODE_DECRYPT, iv, (short) 0, (short) iv.length);

        byte[] dec = new byte[pt.length + 16];
        short decLen = jcGcm.doFinal(out, (short) 0, (short) out.length, dec, (short) 0);
        assertEquals(pt.length, decLen, "jcardsim GCM decrypt returns pt length");

        byte[] ptOut = new byte[pt.length];
        System.arraycopy(dec, 0, ptOut, 0, pt.length);
        assertArrayEquals(pt, ptOut,
            "AliroGcm ciphertext must decrypt correctly under jcardsim AEADCipher");
    }

    /**
     * Partial final block: plaintext whose length isn't a multiple of 16.
     * GCTR must handle the partial block correctly (XOR only |P| mod 16 bytes
     * of the last keystream block), and GHASH must zero-pad.
     */
    @Test
    void partialFinalBlock() {
        byte[] key = new byte[32];
        for (int i = 0; i < 32; i++) key[i] = (byte) i;
        byte[] iv = new byte[12];
        iv[11] = 0x01;
        byte[] pt = new byte[17];  // 1 full block + 1 byte
        for (int i = 0; i < 17; i++) pt[i] = (byte) (0xAA + i);

        byte[] out = new byte[pt.length + 16];
        AliroGcm gcm = new AliroGcm();
        short n = gcm.encrypt(key, (short) 0, iv, (short) 0,
            pt, (short) 0, (short) pt.length, out, (short) 0);

        assertEquals(pt.length + 16, n,
            "ciphertext length must equal plaintext length (GCM is length-preserving)");

        // Round-trip through jcardsim for correctness.
        AESKey k = (AESKey) KeyBuilder.buildKey(
            KeyBuilder.TYPE_AES, KeyBuilder.LENGTH_AES_256, false);
        k.setKey(key, (short) 0);
        AEADCipher jcGcm = (AEADCipher) Cipher.getInstance(AEADCipher.ALG_AES_GCM, false);
        jcGcm.init(k, Cipher.MODE_DECRYPT, iv, (short) 0, (short) iv.length);
        byte[] dec = new byte[pt.length + 16];
        short decLen = jcGcm.doFinal(out, (short) 0, (short) out.length, dec, (short) 0);
        assertEquals(pt.length, decLen);
        byte[] ptOut = new byte[pt.length];
        System.arraycopy(dec, 0, ptOut, 0, pt.length);
        assertArrayEquals(pt, ptOut, "17-byte plaintext must round-trip");
    }

    @Test
    void differentKeysProduceDifferentTags() {
        byte[] iv = new byte[12];
        iv[11] = 0x01;
        byte[] pt = new byte[] { 0x11, 0x22, 0x33, 0x44 };

        byte[] key1 = new byte[32];
        byte[] key2 = new byte[32];
        key2[0] = 0x01;  // differ in one bit

        AliroGcm gcm = new AliroGcm();
        byte[] out1 = new byte[pt.length + 16];
        byte[] out2 = new byte[pt.length + 16];
        gcm.encrypt(key1, (short) 0, iv, (short) 0, pt, (short) 0, (short) pt.length, out1, (short) 0);
        gcm.encrypt(key2, (short) 0, iv, (short) 0, pt, (short) 0, (short) pt.length, out2, (short) 0);

        // Extract tags (last 16 bytes).
        byte[] tag1 = new byte[16];
        byte[] tag2 = new byte[16];
        System.arraycopy(out1, pt.length, tag1, 0, 16);
        System.arraycopy(out2, pt.length, tag2, 0, 16);

        assertFalse(java.util.Arrays.equals(tag1, tag2),
            "different keys must produce different tags");
    }

    @Test
    void differentIvsProduceDifferentCiphertexts() {
        byte[] key = new byte[32];
        for (int i = 0; i < 32; i++) key[i] = (byte) i;
        byte[] pt = new byte[32];
        for (int i = 0; i < 32; i++) pt[i] = (byte) (0x10 + i);

        byte[] iv1 = new byte[12];
        iv1[11] = 0x01;
        byte[] iv2 = new byte[12];
        iv2[11] = 0x02;

        AliroGcm gcm = new AliroGcm();
        byte[] out1 = new byte[pt.length + 16];
        byte[] out2 = new byte[pt.length + 16];
        gcm.encrypt(key, (short) 0, iv1, (short) 0, pt, (short) 0, (short) pt.length, out1, (short) 0);
        gcm.encrypt(key, (short) 0, iv2, (short) 0, pt, (short) 0, (short) pt.length, out2, (short) 0);

        byte[] ct1 = new byte[pt.length];
        byte[] ct2 = new byte[pt.length];
        System.arraycopy(out1, 0, ct1, 0, pt.length);
        System.arraycopy(out2, 0, ct2, 0, pt.length);

        assertFalse(java.util.Arrays.equals(ct1, ct2),
            "different IVs must produce different ciphertexts (CTR stream differs)");
    }

    private static byte[] hex(String s) {
        s = s.replaceAll("\\s+", "");
        byte[] out = new byte[s.length() / 2];
        for (int i = 0; i < out.length; i++) {
            out[i] = (byte) Integer.parseInt(s.substring(i * 2, i * 2 + 2), 16);
        }
        return out;
    }
}