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");
    }

    /**
     * Rekey-across-encrypts: M2D.3's stream-encrypt path calls
     * {@link AliroGcm#encrypt} twice in a single session (once for the
     * EXCHANGE Reader Status ack, once for the ENVELOPE DeviceResponse) with
     * a different (key, IV) pair each time. This test pins that no leftover
     * state from the first call (cached H, M-table, AES key slot) corrupts
     * the second. M2C.2 extracts the rekey path into {@code setKeyAndIv};
     * this test is the regression spec for that refactor.
     */
    @Test
    void setKeyAndIvAllowsRekeyAcrossMultipleEncrypts() {
        AliroGcm gcm = new AliroGcm();

        // Encrypt block A under key1 + iv1
        byte[] keyA = makeKey((byte) 0xA0);
        byte[] ivA  = makeIv((byte) 0xAA);
        byte[] ptA  = bytes("hello A");
        byte[] outA = new byte[ptA.length + 16];
        gcm.encrypt(keyA, (short) 0, ivA, (short) 0,
            ptA, (short) 0, (short) ptA.length, outA, (short) 0);

        // Rekey + encrypt block B under key2 + iv2
        byte[] keyB = makeKey((byte) 0xB0);
        byte[] ivB  = makeIv((byte) 0xBB);
        byte[] ptB  = bytes("hello B");
        byte[] outB = new byte[ptB.length + 16];
        gcm.encrypt(keyB, (short) 0, ivB, (short) 0,
            ptB, (short) 0, (short) ptB.length, outB, (short) 0);

        // Both must round-trip cleanly through decrypt under their own key+iv
        byte[] ptAroundTrip = new byte[ptA.length];
        gcm.decrypt(keyA, (short) 0, ivA, (short) 0,
            outA, (short) 0, (short) outA.length, ptAroundTrip, (short) 0);
        assertArrayEquals(ptA, ptAroundTrip,
            "block A must decrypt under (keyA, ivA) after gcm has been re-keyed to B");

        byte[] ptBroundTrip = new byte[ptB.length];
        gcm.decrypt(keyB, (short) 0, ivB, (short) 0,
            outB, (short) 0, (short) outB.length, ptBroundTrip, (short) 0);
        assertArrayEquals(ptB, ptBroundTrip,
            "block B must decrypt under (keyB, ivB) — no leftover state from A");
    }

    private static byte[] makeKey(byte seed) {
        byte[] k = new byte[32];
        for (int i = 0; i < 32; i++) k[i] = (byte) (seed + i);
        return k;
    }

    private static byte[] makeIv(byte seed) {
        byte[] iv = new byte[12];
        for (int i = 0; i < 12; i++) iv[i] = (byte) (seed + i);
        return iv;
    }

    private static byte[] bytes(String s) {
        return s.getBytes(java.nio.charset.StandardCharsets.UTF_8);
    }

    @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;
    }
}
