package com.dangerousthings.aliro;

import com.licel.jcardsim.smartcardio.CardSimulator;
import java.security.KeyPair;
import java.security.SecureRandom;
import java.security.interfaces.ECPublicKey;
import javacard.framework.AID;
import javax.smartcardio.CommandAPDU;
import javax.smartcardio.ResponseAPDU;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotEquals;

/**
 * AUTH1 tests. Covers state-machine sequencing, input TLV validation, and
 * provisioning preconditions. Real cryptographic verification of the
 * reader signature and construction of the encrypted response come in
 * later iterations.
 */
class AliroAppletAuth1Test {

    private static final byte INS_AUTH1 = (byte) 0x81;

    private CardSimulator sim;
    private KeyPair credentialKeyPair;
    private ReaderSide reader;
    private byte[] credentialEphemPubKey;  // captured from the last AUTH0 response

    @BeforeEach
    void setUp() {
        // reader.provision() below installs PersonalizationApplet, whose
        // constructor calls CredentialStore.bootstrap() and gives us a fresh
        // store for every test — no explicit reset needed any more.
        sim = new CardSimulator();
        AID expeditedAid = new AID(AliroAids.EXPEDITED, (short) 0, (byte) AliroAids.EXPEDITED.length);
        sim.installApplet(expeditedAid, AliroApplet.class);

        credentialKeyPair = Auth0Command.generateEphemeralKeyPair();
        reader = new ReaderSide();
        reader.provision(sim, credentialKeyPair);

        selectExpedited();
    }

    private void selectExpedited() {
        CommandAPDU select = new CommandAPDU(0x00, 0xA4, 0x04, 0x00, AliroAids.EXPEDITED, 256);
        ResponseAPDU r = sim.transmitCommand(select);
        assertEquals(0x9000, r.getSW(), "SELECT must succeed");
    }

    /** Runs a standard-phase AUTH0 and captures the returned credential_ePubK for later AUTH1 use. */
    private void sendStandardAuth0() {
        reader.startTransaction();
        byte[] data = reader.buildAuth0Data();
        CommandAPDU auth0 = new CommandAPDU(
            Auth0Command.CLA & 0xFF, Auth0Command.INS & 0xFF, 0x00, 0x00, data, 256);
        ResponseAPDU r = sim.transmitCommand(auth0);
        assertEquals(0x9000, r.getSW(), "AUTH0 must succeed as a precondition for AUTH1 tests");
        credentialEphemPubKey = TlvUtil.findTopLevel(r.getData(), 0x86);
    }

    /** AUTH1 command with a correctly-signed reader signature. */
    private ResponseAPDU sendValidAuth1() {
        return sendAuth1(reader.buildAuth1Data(credentialEphemPubKey));
    }

    /** Syntactically-valid AUTH1 payload with zeroed signature (for TLV-only tests). */
    private ResponseAPDU sendStubAuth1() {
        byte[] data = new byte[3 + 2 + 64];
        int i = 0;
        data[i++] = 0x41; data[i++] = 0x01; data[i++] = 0x00;
        data[i++] = (byte) 0x9E; data[i++] = 0x40;
        return sendAuth1(data);
    }

    private ResponseAPDU sendAuth1(byte[] data) {
        CommandAPDU auth1 = new CommandAPDU(
            Auth0Command.CLA & 0xFF, INS_AUTH1 & 0xFF, 0x00, 0x00, data, 256);
        return sim.transmitCommand(auth1);
    }

    @Test
    void auth1WithoutPriorAuth0IsRejected() {
        ResponseAPDU r = sendStubAuth1();
        assertNotEquals(0x9000, r.getSW(),
            "AUTH1 must fail unless immediately preceded by a successful AUTH0");
        assertEquals(0x6985, r.getSW(),
            "expected SW_CONDITIONS_NOT_SATISFIED (0x6985) when AUTH0 not done");
    }

    @Test
    void auth1WithValidReaderSignatureIsAccepted() {
        sendStandardAuth0();
        ResponseAPDU r = sendValidAuth1();
        assertEquals(0x9000, r.getSW(),
            "AUTH1 with a correctly-signed reader signature must succeed");
    }

    @Test
    void auth1ReplayWithoutFreshAuth0IsRejected() {
        sendStandardAuth0();
        ResponseAPDU first = sendValidAuth1();
        assertEquals(0x9000, first.getSW());

        // Second AUTH1 without a new AUTH0 must fail — otherwise the card would
        // re-encrypt with device_counter=1 under the same key, reusing the GCM
        // IV and catastrophically breaking confidentiality + integrity.
        ResponseAPDU replay = sendValidAuth1();
        assertEquals(0x6985, replay.getSW(),
            "AUTH1 after a prior successful AUTH1 must require a fresh AUTH0");
    }

    @Test
    void auth1SuccessArmsStepUpContextWithExpectedStepUpSK() throws Exception {
        sendStandardAuth0();
        ResponseAPDU r = sendValidAuth1();
        assertEquals(0x9000, r.getSW());

        // Reader derives the full expedited-standard 160B key material and
        // extracts StepUpSK (offset 64..96, per spec §8.3.1.13).
        byte[] expected = java.util.Arrays.copyOfRange(
            reader.deriveExpeditedKeyMaterial(credentialEphemPubKey), 64, 96);

        org.junit.jupiter.api.Assertions.assertTrue(SessionContext.isStepUpArmed(),
            "successful AUTH1 must arm the step-up context");
        byte[] got = new byte[32];
        SessionContext.copyStepUpSKForTesting(got, (short) 0);
        org.junit.jupiter.api.Assertions.assertArrayEquals(expected, got,
            "armed StepUpSK must equal derived_keys_volatile[64..96]");
    }

    @Test
    void freshAuth0DisarmsPreviousStepUpArming() throws Exception {
        // First transaction arms the step-up context.
        sendStandardAuth0();
        assertEquals(0x9000, sendValidAuth1().getSW());
        org.junit.jupiter.api.Assertions.assertTrue(SessionContext.isStepUpArmed());

        // Starting a new transaction via AUTH0 (without re-SELECT) must
        // invalidate the prior arming — otherwise a mid-transaction failure
        // could leave stale StepUpSK reachable by a subsequent SELECT STEP_UP.
        reader.startTransaction();
        byte[] data = reader.buildAuth0Data();
        CommandAPDU auth0 = new CommandAPDU(
            Auth0Command.CLA & 0xFF, Auth0Command.INS & 0xFF, 0x00, 0x00, data, 256);
        assertEquals(0x9000, sim.transmitCommand(auth0).getSW());
        org.junit.jupiter.api.Assertions.assertFalse(SessionContext.isStepUpArmed(),
            "starting a new AUTH0 must disarm any prior step-up state");
    }

    @Test
    void reSelectExpeditedDisarmsStepUpContext() throws Exception {
        sendStandardAuth0();
        assertEquals(0x9000, sendValidAuth1().getSW());
        org.junit.jupiter.api.Assertions.assertTrue(SessionContext.isStepUpArmed());

        selectExpedited();
        org.junit.jupiter.api.Assertions.assertFalse(SessionContext.isStepUpArmed(),
            "re-SELECT of expedited AID starts a new transaction and must disarm step-up");
    }

    @Test
    void auth1_bitmap_adNotProvisioned_returns0x0004() throws Exception {
        // Default setUp provisions keys but no Access Document. Bit 2
        // (step-up AID SELECT required) is still set — the split-AID step-up
        // architecture is always advertised, even when there's nothing
        // behind it.
        sendStandardAuth0();
        ResponseAPDU r = sendValidAuth1();
        assertEquals(0x9000, r.getSW());

        byte[] pt = ReaderSide.decryptAuth1Response(
            reader.deriveExpeditedSKDevice(credentialEphemPubKey), r.getData());
        byte[] bitmap = TlvUtil.findTopLevel(pt, 0x5E);
        org.junit.jupiter.api.Assertions.assertArrayEquals(
            new byte[] { 0x00, 0x04 }, bitmap,
            "no Access Document → bit 2 only (0x0004); bit 0 reflects actual retrievability");
    }

    @Test
    void auth1_bitmap_adProvisionedButNotVerified_returns0x0004() throws Exception {
        // Provision opaque AD bytes and mark them finalized, but NOT verified.
        // Post-M2A.3 the finalize flow requires IssuerAuth verify so this is a
        // defensive case — bit 0 (AD retrievable) must NOT be advertised since
        // we wouldn't actually serve unverified AD on the step-up path.
        byte[] ad = new byte[120];
        for (int i = 0; i < ad.length; i++) ad[i] = (byte) (i ^ 0x5A);
        org.junit.jupiter.api.Assertions.assertTrue(
            CredentialStore.get().writeAccessDocumentChunk(ad, (short) 0, (short) 0, (short) ad.length));
        CredentialStore.get().markAccessDocumentFinalizedForTesting((short) ad.length);
        // Deliberately do NOT mark verified.

        sendStandardAuth0();
        ResponseAPDU r = sendValidAuth1();
        assertEquals(0x9000, r.getSW());

        byte[] pt = ReaderSide.decryptAuth1Response(
            reader.deriveExpeditedSKDevice(credentialEphemPubKey), r.getData());
        byte[] bitmap = TlvUtil.findTopLevel(pt, 0x5E);
        org.junit.jupiter.api.Assertions.assertArrayEquals(
            new byte[] { 0x00, 0x04 }, bitmap,
            "AD finalized but not verified → bit 2 only (0x0004); bit 0 requires verified flag");
    }

    @Test
    void auth1_bitmap_adProvisionedAndVerified_returns0x0005() throws Exception {
        // Provision a non-empty Access Document via CredentialStore (bypasses
        // the PersonalizationApplet INS layer — that flow is tested
        // separately). Mark both finalized AND verified to mirror the
        // post-M2A.3 happy path.
        byte[] ad = new byte[120];
        for (int i = 0; i < ad.length; i++) ad[i] = (byte) (i ^ 0x5A);
        org.junit.jupiter.api.Assertions.assertTrue(
            CredentialStore.get().writeAccessDocumentChunk(ad, (short) 0, (short) 0, (short) ad.length));
        CredentialStore.get().markAccessDocumentFinalizedForTesting((short) ad.length);
        CredentialStore.get().markAccessDocumentVerifiedForTesting();

        sendStandardAuth0();
        ResponseAPDU r = sendValidAuth1();
        assertEquals(0x9000, r.getSW());

        byte[] pt = ReaderSide.decryptAuth1Response(
            reader.deriveExpeditedSKDevice(credentialEphemPubKey), r.getData());
        byte[] bitmap = TlvUtil.findTopLevel(pt, 0x5E);
        // bit 0 (access doc retrievable) + bit 2 (requires step-up AID on NFC)
        //   = 2^0 + 2^2 = 0x0005 (big-endian)
        org.junit.jupiter.api.Assertions.assertArrayEquals(
            new byte[] { 0x00, 0x05 }, bitmap,
            "AD finalized AND verified → bits 0 + 2 set (0x0005)");
    }

    @Test
    void auth1WithKeySlotRequestReturnsSha1KeyIdentifier() throws Exception {
        sendStandardAuth0();
        // Build a valid AUTH1 payload, then flip command_parameters bit 0 to 0
        // (key_slot path). The reader sig is computed over Table 8-12, which
        // doesn't include command_parameters, so the sig stays valid.
        byte[] data = reader.buildAuth1Data(credentialEphemPubKey);
        data[2] = 0x00;
        ResponseAPDU r = sendAuth1(data);
        assertEquals(0x9000, r.getSW());

        byte[] pt = ReaderSide.decryptAuth1Response(
            reader.deriveExpeditedSKDevice(credentialEphemPubKey, (byte) 0x00), r.getData());

        byte[] keySlot = TlvUtil.findTopLevel(pt, 0x4E);
        byte[] credPubTag = TlvUtil.findTopLevel(pt, 0x5A);
        org.junit.jupiter.api.Assertions.assertNotNull(keySlot, "0x4E key_slot must be present when cmd_params bit 0 = 0");
        org.junit.jupiter.api.Assertions.assertNull(credPubTag, "0x5A credential_PubK must NOT be present when cmd_params bit 0 = 0");
        assertEquals(8, keySlot.length, "key_slot is the first 8 bytes of keyIdentifier");

        // Expected: first 8 bytes of SHA-1(uncompressed credential long-term pubkey)
        java.security.MessageDigest sha1 = java.security.MessageDigest.getInstance("SHA-1");
        byte[] credPubUncomp = new byte[65];
        credPubUncomp[0] = 0x04;
        byte[] x = toFixed32(credentialKeyPair.getPublic() instanceof java.security.interfaces.ECPublicKey
            ? ((java.security.interfaces.ECPublicKey) credentialKeyPair.getPublic()).getW().getAffineX().toByteArray()
            : new byte[0]);
        byte[] y = toFixed32(((java.security.interfaces.ECPublicKey) credentialKeyPair.getPublic())
            .getW().getAffineY().toByteArray());
        System.arraycopy(x, 0, credPubUncomp, 1, 32);
        System.arraycopy(y, 0, credPubUncomp, 33, 32);
        byte[] expected = java.util.Arrays.copyOf(sha1.digest(credPubUncomp), 8);
        org.junit.jupiter.api.Assertions.assertArrayEquals(expected, keySlot,
            "key_slot must be SHA-1(uncompressed credential_PubK)[0:8]");
    }

    private static byte[] toFixed32(byte[] in) {
        if (in.length == 32) return in;
        byte[] out = new byte[32];
        if (in.length > 32) System.arraycopy(in, in.length - 32, out, 0, 32);
        else System.arraycopy(in, 0, out, 32 - in.length, in.length);
        return out;
    }

    @Test
    void distinctTransactionsProduceDistinctCiphertexts() throws Exception {
        sendStandardAuth0();
        ResponseAPDU r1 = sendValidAuth1();
        assertEquals(0x9000, r1.getSW());
        byte[] ct1 = r1.getData();
        byte[] plaintext1 = ReaderSide.decryptAuth1Response(
            reader.deriveExpeditedSKDevice(credentialEphemPubKey), ct1);

        // Second full transaction: fresh AUTH0 → fresh ephemerals + txn_id.
        sendStandardAuth0();
        ResponseAPDU r2 = sendValidAuth1();
        assertEquals(0x9000, r2.getSW());
        byte[] ct2 = r2.getData();
        byte[] plaintext2 = ReaderSide.decryptAuth1Response(
            reader.deriveExpeditedSKDevice(credentialEphemPubKey), ct2);

        org.junit.jupiter.api.Assertions.assertFalse(
            java.util.Arrays.equals(ct1, ct2),
            "two transactions must produce different ciphertexts — catches hardcoded salt/IV/key bugs");

        // Plaintexts also differ (UD sig is over Table 8-13 which includes txn_id).
        byte[] sig1 = TlvUtil.findTopLevel(plaintext1, 0x9E);
        byte[] sig2 = TlvUtil.findTopLevel(plaintext2, 0x9E);
        org.junit.jupiter.api.Assertions.assertFalse(
            java.util.Arrays.equals(sig1, sig2),
            "UD signatures over Table 8-13 must differ across transactions");
    }

    @Test
    void auth1EncryptedResponseDecryptsAndVerifies() throws Exception {
        sendStandardAuth0();
        ResponseAPDU r = sendValidAuth1();
        assertEquals(0x9000, r.getSW());

        byte[] ctAndTag = r.getData();
        byte[] skDevice = reader.deriveExpeditedSKDevice(credentialEphemPubKey);
        byte[] plaintext = ReaderSide.decryptAuth1Response(skDevice, ctAndTag);

        byte[] credPub = TlvUtil.findTopLevel(plaintext, 0x5A);
        byte[] udSig   = TlvUtil.findTopLevel(plaintext, 0x9E);
        byte[] bitmap  = TlvUtil.findTopLevel(plaintext, 0x5E);

        org.junit.jupiter.api.Assertions.assertNotNull(credPub, "Table 8-11 must include 0x5A (credential_PubK)");
        org.junit.jupiter.api.Assertions.assertNotNull(udSig,   "Table 8-11 must include 0x9E (UD signature)");
        org.junit.jupiter.api.Assertions.assertNotNull(bitmap,  "Table 8-11 must include 0x5E (signaling_bitmap)");
        assertEquals(65, credPub.length);
        assertEquals(64, udSig.length);
        assertEquals(2,  bitmap.length);

        byte[] table813 = reader.buildTable813(credentialEphemPubKey);
        org.junit.jupiter.api.Assertions.assertTrue(
            ReaderSide.verifyUdSig(credPub, table813, udSig),
            "UD signature in the decrypted Table 8-11 must verify under credential_PubK over Table 8-13");
    }

    @Test
    void auth1WithInvalidReaderSignatureFails() {
        sendStandardAuth0();
        // Valid TLV structure but a random (definitely-not-matching) sig.
        byte[] data = new byte[3 + 2 + 64];
        int i = 0;
        data[i++] = 0x41; data[i++] = 0x01; data[i++] = 0x00;
        data[i++] = (byte) 0x9E; data[i++] = 0x40;
        new SecureRandom().nextBytes(java.util.Arrays.copyOfRange(data, i, i + 64));
        // Force at least one non-zero byte so the sig isn't accidentally valid.
        data[i] = 0x11;
        data[i + 1] = 0x22;
        ResponseAPDU r = sendAuth1(data);
        assertNotEquals(0x9000, r.getSW(),
            "AUTH1 with a bad reader signature must not return success");
    }

    @Test
    void reSelectClearsAuth0State() {
        sendStandardAuth0();
        selectExpedited();
        ResponseAPDU r = sendStubAuth1();
        assertEquals(0x6985, r.getSW(),
            "re-SELECT must invalidate previous AUTH0 state");
    }

    @Test
    void auth1MissingCommandParametersFails() {
        sendStandardAuth0();
        byte[] data = new byte[2 + 64];
        data[0] = (byte) 0x9E; data[1] = 0x40;
        assertEquals(0x6A80, sendAuth1(data).getSW(),
            "missing tag 0x41 must be rejected with SW_WRONG_DATA");
    }

    @Test
    void auth1MissingReaderSignatureFails() {
        sendStandardAuth0();
        byte[] data = new byte[] { 0x41, 0x01, 0x00 };
        assertEquals(0x6A80, sendAuth1(data).getSW(),
            "missing tag 0x9E (reader signature) must be rejected");
    }

    @Test
    void auth1WithWrongReaderSignatureLengthFails() {
        sendStandardAuth0();
        byte[] data = new byte[3 + 2 + 63];
        int i = 0;
        data[i++] = 0x41; data[i++] = 0x01; data[i++] = 0x00;
        data[i++] = (byte) 0x9E; data[i++] = 0x3F;
        assertEquals(0x6A80, sendAuth1(data).getSW(),
            "reader signature length != 64 must be rejected");
    }

    @Test
    void auth1WithReaderCertIsRejectedInV1() {
        sendStandardAuth0();
        byte[] data = new byte[3 + 2 + 64 + 2 + 1];
        int i = 0;
        data[i++] = 0x41; data[i++] = 0x01; data[i++] = 0x00;
        data[i++] = (byte) 0x9E; data[i++] = 0x40;
        i += 64;
        data[i++] = (byte) 0x90; data[i++] = 0x01; data[i++] = (byte) 0xFF;
        assertEquals(0x6A80, sendAuth1(data).getSW(),
            "reader_Cert (tag 0x90) is unsupported in v1 and must be rejected");
    }

    @Test
    void auth1WithUnknownTrailingTagIsAccepted() {
        sendStandardAuth0();
        // Build a real-sig AUTH1 payload, then append an unknown TLV.
        byte[] valid = reader.buildAuth1Data(credentialEphemPubKey);
        byte[] data = new byte[valid.length + 3];
        System.arraycopy(valid, 0, data, 0, valid.length);
        int i = valid.length;
        data[i++] = (byte) 0x7F; data[i++] = 0x01; data[i++] = (byte) 0xAB;
        assertEquals(0x9000, sendAuth1(data).getSW(),
            "unknown trailing TLVs must be accepted (spec §1.4.3)");
    }

    @Test
    void auth1WithoutProvisionedKeysFails() {
        CredentialStore.get().resetForTesting();
        selectExpedited();
        sendStandardAuth0();
        ResponseAPDU r = sendStubAuth1();
        assertEquals(0x6985, r.getSW(),
            "AUTH1 must fail if CredentialStore was not provisioned (6985)");
    }
}