package com.dangerousthings.aliro;

import javacard.security.ECPrivateKey;
import javacard.security.KeyAgreement;
import javacard.security.MessageDigest;

/**
 * Low-level Aliro cryptographic primitives, factored out so they can be
 * unit-tested independently of the APDU dispatch in {@link AliroApplet}.
 *
 * <p>Intended to be instantiated once per applet (or test) and reused —
 * the bundled {@link KeyAgreement} and {@link AliroHmac} objects are
 * expensive to allocate.
 *
 * <p>HMAC-SHA-256 is supplied by {@link AliroHmac}, our userland
 * implementation. The target card (NXP J3R180) does not expose
 * {@code KeyBuilder.TYPE_HMAC} at any tested key size nor
 * {@code Signature.ALG_HMAC_SHA_256}, so we build HMAC on top of the card's
 * raw SHA-256 primitive per RFC 2104.
 */
final class AliroCrypto {

    private static final short HASH_LEN = 32;

    /** Maximum expanded-input buffer size for HKDF-Expand: one T(i-1) block
     *  (32B) + info (up to ~200B, defensively) + one counter byte. Rounded up
     *  to 256B. The largest info we currently pass is the 32-byte salt_volatile
     *  info for expedited-standard derivation; step-up uses 8-byte ASCII info. */
    private static final short EXPAND_SCRATCH_LEN = 256;

    private KeyAgreement ecdhPlain;
    private AliroHmac aliroHmac;

    /** Native SHA-256 instance used by {@link #deriveKdh} (X9.63 KDF). */
    private MessageDigest sha256;

    /** Reusable scratch for one HMAC output (T(i)) and one counter byte. */
    private byte[] hkdfPrevT;

    /** Scratch for assembling the HMAC message for HKDF-Expand: T(i-1) || info || counter. */
    private byte[] expandScratch;

    /** Scratch for {@link #deriveKdh}: 32B for Z_AB then 32B for PRK. */
    private byte[] kdfWorkbuf;
    private static final short KDF_WORKBUF_LEN = 64;

    /** 32 zero bytes — RFC 5869 §2.2 fallback when {@link #hkdfExtract} is
     *  called with an empty salt. Allocated in the constructor; Java Card
     *  bans {@code new} in static initializers, so we can't make this static-final. */
    private byte[] emptySaltZeros;

    AliroCrypto() {
        try {
            ecdhPlain = KeyAgreement.getInstance(KeyAgreement.ALG_EC_SVDP_DH_PLAIN, false);
        } catch (Throwable t) {
            javacard.framework.ISOException.throwIt((short) 0x6FC1);
        }
        // Userland HMAC-SHA-256 on top of MessageDigest.ALG_SHA_256. Diagnostic
        // status word 0x6FC7 gives us one regression cycle of visibility on the
        // real card before we strip it.
        try {
            aliroHmac = new AliroHmac();
        } catch (Throwable t) {
            javacard.framework.ISOException.throwIt((short) 0x6FC7);
        }
        try {
            sha256 = MessageDigest.getInstance(MessageDigest.ALG_SHA_256, false);
        } catch (Throwable t) {
            javacard.framework.ISOException.throwIt((short) 0x6FC2);
        }
        try {
            hkdfPrevT = javacard.framework.JCSystem.makeTransientByteArray(
                HASH_LEN, javacard.framework.JCSystem.CLEAR_ON_DESELECT);
            expandScratch = javacard.framework.JCSystem.makeTransientByteArray(
                EXPAND_SCRATCH_LEN, javacard.framework.JCSystem.CLEAR_ON_DESELECT);
            kdfWorkbuf = javacard.framework.JCSystem.makeTransientByteArray(
                KDF_WORKBUF_LEN, javacard.framework.JCSystem.CLEAR_ON_DESELECT);
        } catch (Throwable t) {
            javacard.framework.ISOException.throwIt((short) 0x6FC4);
        }
        try {
            emptySaltZeros = new byte[32];  // persistent, defaults to all zeros
        } catch (Throwable t) {
            javacard.framework.ISOException.throwIt((short) 0x6FC5);
        }
    }

    /**
     * Computes the x-coordinate of the ECDH shared point, given a local
     * P-256 private key and the peer's uncompressed public point
     * (0x04 || x || y, 65 bytes). Returns 32 (the x-coordinate length).
     */
    short computeEcdhSharedX(
            ECPrivateKey priv,
            byte[] peerPubUncomp, short peerPubOff,
            byte[] out, short outOff) {
        ecdhPlain.init(priv);
        return ecdhPlain.generateSecret(peerPubUncomp, peerPubOff, (short) 65, out, outOff);
    }

    /**
     * Diagnostic-only: exposes the underlying HMAC-SHA-256 primitive so
     * AliroApplet's INS_DIAG_HMAC handler can profile per-call cost. Not
     * used by production AUTH0/AUTH1 paths (those go through hkdfExtract
     * and hkdfExpand).
     */
    short diagHmac(
            byte[] key, short keyOff, short keyLen,
            byte[] msg, short msgOff, short msgLen,
            byte[] out, short outOff) {
        return aliroHmac.compute(key, keyOff, keyLen, msg, msgOff, msgLen, out, outOff);
    }

    /**
     * HKDF-Extract per RFC 5869 §2.2 with HMAC-SHA-256:
     * {@code PRK = HMAC(salt, IKM)}.
     * Returns 32 (HashLen).
     */
    short hkdfExtract(
            byte[] salt, short saltOff, short saltLen,
            byte[] ikm, short ikmOff, short ikmLen,
            byte[] prk, short prkOff) {
        // RFC 5869 §2.2: empty salt is treated as HashLen zero bytes.
        if (saltLen == 0) {
            return aliroHmac.compute(
                emptySaltZeros, (short) 0, HASH_LEN,
                ikm, ikmOff, ikmLen,
                prk, prkOff);
        }
        return aliroHmac.compute(
            salt, saltOff, saltLen,
            ikm, ikmOff, ikmLen,
            prk, prkOff);
    }

    /**
     * HKDF-Expand per RFC 5869 §2.3 with HMAC-SHA-256. Generates {@code L}
     * bytes of output keying material, where {@code L} must be ≤ 255 *
     * HashLen (8160 bytes — far above any Aliro need). The info region must
     * be at most {@code EXPAND_SCRATCH_LEN - HASH_LEN - 1} bytes long
     * (= 223 bytes with the current buffer sizing).
     */
    short hkdfExpand(
            byte[] prk, short prkOff, short prkLen,
            byte[] info, short infoOff, short infoLen,
            short L,
            byte[] out, short outOff) {

        // T(0) = empty; T(i) = HMAC(PRK, T(i-1) || info || i). Concatenate
        // T(1)..T(N) and truncate to L bytes.
        short produced = 0;
        short prevTLen = 0;
        byte counter = (byte) 1;

        while (produced < L) {
            // Assemble the HMAC message as T(i-1) || info || counter in expandScratch.
            short msgLen = 0;
            if (prevTLen != 0) {
                javacard.framework.Util.arrayCopyNonAtomic(
                    hkdfPrevT, (short) 0, expandScratch, msgLen, prevTLen);
                msgLen += prevTLen;
            }
            if (infoLen != 0) {
                javacard.framework.Util.arrayCopyNonAtomic(
                    info, infoOff, expandScratch, msgLen, infoLen);
                msgLen += infoLen;
            }
            expandScratch[msgLen] = counter;
            msgLen++;

            aliroHmac.compute(
                prk, prkOff, prkLen,
                expandScratch, (short) 0, msgLen,
                hkdfPrevT, (short) 0);
            prevTLen = HASH_LEN;

            short remaining = (short) (L - produced);
            short toCopy = remaining < HASH_LEN ? remaining : HASH_LEN;
            javacard.framework.Util.arrayCopyNonAtomic(
                hkdfPrevT, (short) 0, out, (short) (outOff + produced), toCopy);
            produced += toCopy;
            counter++;
        }

        // Wipe the assembled message scratch — it held T(i-1) values (derived
        // from the HKDF PRK) and the info bytes. CLEAR_ON_DESELECT cleans up
        // at deselect, but multi-APDU sessions would otherwise see stale T(i).
        javacard.framework.Util.arrayFillNonAtomic(
            expandScratch, (short) 0, EXPAND_SCRATCH_LEN, (byte) 0);

        return produced;
    }

    /**
     * Derives Kdh per Aliro §8.3.1.4: <b>X9.63 KDF</b> from BSI TR-03111 with
     * H=SHA-256, ZAB = ECDH shared-secret x-coord, SharedInfo =
     * transaction_identifier, K = 256 bits. For 32-byte output X9.63 KDF
     * reduces to a single SHA-256 invocation:
     * <pre>
     *   Kdh = SHA-256(ZAB || 0x00000001 || transaction_identifier)
     * </pre>
     *
     * <p><b>History:</b> we previously misread the §8.3.1.4 note ("actual key
     * derivation is performed using §8.3.1.5") as authorizing HKDF
     * substitution for Kdh itself. The note is about the subsequent
     * session-key derivation (§8.3.1.13 uses §8.3.1.5 HKDF), not Kdh. The
     * misread was symmetric between this applet and our PC/SC reader, so
     * AUTH1 succeeded against our own host-side reader but failed against
     * ST's X-CUBE-ALIRO library (which follows §8.3.1.4 correctly) with
     * {@code ACWG_Error_Crypto_EncryptDecrypt}. Fixed 2026-06-11.
     *
     * <p>Writes 32 bytes to {@code out[outOff..]} and returns 32.
     */
    short deriveKdh(
            ECPrivateKey priv,
            byte[] peerPubUncomp, short peerPubOff,
            byte[] txnId, short txnIdOff, short txnIdLen,
            byte[] out, short outOff) {
        // Stage ZAB || counter(0x00000001) at kdfWorkbuf[0..36).
        computeEcdhSharedX(priv, peerPubUncomp, peerPubOff, kdfWorkbuf, (short) 0);
        kdfWorkbuf[32] = 0;
        kdfWorkbuf[33] = 0;
        kdfWorkbuf[34] = 0;
        kdfWorkbuf[35] = 1;
        // Kdh = SHA-256(ZAB || counter || txnId) -- one shot.
        sha256.reset();
        sha256.update(kdfWorkbuf, (short) 0, (short) 36);
        sha256.doFinal(txnId, txnIdOff, txnIdLen, out, outOff);
        // Wipe ZAB from working memory.
        javacard.framework.Util.arrayFillNonAtomic(kdfWorkbuf, (short) 0, (short) 36, (byte) 0);
        return HASH_LEN;
    }

    /** ASCII info strings for step-up session key derivation, per mdoc 9.1.1.5. */
    private static final byte[] STEP_UP_INFO_SK_DEVICE = {
        (byte) 'S', (byte) 'K', (byte) 'D', (byte) 'e',
        (byte) 'v', (byte) 'i', (byte) 'c', (byte) 'e'
    };
    private static final byte[] STEP_UP_INFO_SK_READER = {
        (byte) 'S', (byte) 'K', (byte) 'R', (byte) 'e',
        (byte) 'a', (byte) 'd', (byte) 'e', (byte) 'r'
    };

    /**
     * Derives {@code StepUpSKDevice} and {@code StepUpSKReader} per Aliro
     * §8.4.3, which points at mdoc [6] clause 9.1.1.5 with two Aliro-specific
     * changes: IKM is the {@code StepUpSK} computed in §8.3.1.13 and the salt
     * is empty. Info remains the ASCII bytes of {@code "SKDevice"} and
     * {@code "SKReader"}.
     *
     * <p>Writes 32 bytes each to {@code skDeviceOut[skDeviceOff..]} and
     * {@code skReaderOut[skReaderOff..]}.
     */
    short deriveStepUpSessionKeys(
            byte[] stepUpSK, short stepUpSKOff,
            byte[] skDeviceOut, short skDeviceOff,
            byte[] skReaderOut, short skReaderOff) {
        // HKDF-Extract once (shared across the two Expand calls): empty salt.
        hkdfExtract(
            kdfWorkbuf, (short) 0, (short) 0,          // empty salt — use first 0 bytes of anything
            stepUpSK, stepUpSKOff, HASH_LEN,
            kdfWorkbuf, HASH_LEN);                     // PRK at [32..64)

        hkdfExpand(
            kdfWorkbuf, HASH_LEN, HASH_LEN,
            STEP_UP_INFO_SK_DEVICE, (short) 0, (short) STEP_UP_INFO_SK_DEVICE.length,
            HASH_LEN,
            skDeviceOut, skDeviceOff);

        hkdfExpand(
            kdfWorkbuf, HASH_LEN, HASH_LEN,
            STEP_UP_INFO_SK_READER, (short) 0, (short) STEP_UP_INFO_SK_READER.length,
            HASH_LEN,
            skReaderOut, skReaderOff);

        return HASH_LEN;
    }

    /**
     * Derives the 160-byte {@code derived_keys_volatile} for the
     * expedited-standard phase per Aliro §8.3.1.13. Output is laid out as
     * ExpeditedSKReader[0..32), ExpeditedSKDevice[32..64), StepUpSK[64..96),
     * BleSK[96..128), URSK[128..160). For NFC-only operation only the
     * first 96 bytes are load-bearing, but the spec mandates generating all
     * 160 so the caller can discard what it doesn't need.
     */
    short deriveExpeditedStandardKeys(
            byte[] kdh, short kdhOff,
            byte[] saltVolatile, short saltOff, short saltLen,
            byte[] info, short infoOff, short infoLen,
            byte[] out, short outOff) {
        hkdfExtract(
            saltVolatile, saltOff, saltLen,
            kdh, kdhOff, HASH_LEN,
            kdfWorkbuf, (short) 0);
        hkdfExpand(
            kdfWorkbuf, (short) 0, HASH_LEN,
            info, infoOff, infoLen,
            (short) 160,
            out, outOff);
        return (short) 160;
    }
}