"""Smoke tests for step-up reader verification.

Covers:
  - StepUpSK -> SKDevice/SKReader HKDF parity with stdlib RFC 5869.
  - M2G.1: `verify_step_up_m2` drives the full M2 ENVELOPE round-trip:
        ENVELOPE(CBOR DeviceRequest) -> 252B chunk + SW=61xx
        GET RESPONSE                 -> remainder + SW=9000
        decrypt(StepUpSKDevice, deviceIv(1)) -> canonical DeviceResponse
        extract documents[0].issuerSigned.issuerAuth, assert == expected_ad.

The M2G.1 test uses a mock transmit driven by a canned ciphertext that the
test itself synthesizes via AESGCM — so it pins the reader's framing /
chaining / decrypt / CBOR-walk logic without depending on the applet build.
M2G.2 covers the real-hardware verdict.
"""

import hashlib
import hmac

import cbor2
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

from aliro_harness.reader.crypto import derive_step_up_session_keys
from aliro_harness.reader.step_up import verify_step_up_m2


def _hkdf_manual(ikm: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    """Bare-stdlib RFC 5869 reference. Independent of `cryptography` so the
    cross-check actually catches a wrong-info / wrong-salt regression in
    derive_step_up_session_keys."""
    use_salt = salt if salt else b"\x00" * 32
    prk = hmac.new(use_salt, ikm, hashlib.sha256).digest()
    out = b""
    t = b""
    counter = 1
    while len(out) < length:
        t = hmac.new(prk, t + info + bytes([counter]), hashlib.sha256).digest()
        out += t
        counter += 1
    return out[:length]


def test_derive_step_up_session_keys_matches_rfc5869():
    # Same IKM as AliroCryptoTest.deriveStepUpSessionKeysMatchesManualHkdf:
    # bytes 0xC0..0xDF -- so this Python test pins the same vector the
    # applet test pins.
    step_up_sk = bytes(range(0xC0, 0xE0))

    sk_device, sk_reader = derive_step_up_session_keys(step_up_sk)

    assert len(sk_device) == 32
    assert len(sk_reader) == 32
    assert sk_device != sk_reader
    assert sk_device != step_up_sk

    assert sk_device == _hkdf_manual(step_up_sk, b"", b"SKDevice", 32)
    assert sk_reader == _hkdf_manual(step_up_sk, b"", b"SKReader", 32)


# ------------------------------------------------------------------------- #
# M2G.1: full ENVELOPE round-trip with GET RESPONSE chaining + AD assert.   #
# ------------------------------------------------------------------------- #


# Same canonical-CBOR shape the applet's DeviceResponseBuilder produces and
# the M2G.1 verify function walks. Used to build a synthetic
# (decrypts-cleanly) ciphertext for the mock transmit.
def _build_fixed_device_response(ad_bytes: bytes) -> bytes:
    """Mirror of applet DeviceResponseBuilder.build for the M2 fixed shape.
    `issuerAuth` is the raw AD bytes spliced verbatim — we decode them so
    cbor2 re-encodes canonically (matches the applet which assumes the AD
    was canonical at personalization)."""
    return cbor2.dumps(
        {
            "status": 0,
            "version": "1.0",
            "documents": [
                {
                    "docType": "org.iso.18013.5.1.mDL",
                    "issuerSigned": {
                        "issuerAuth": cbor2.loads(ad_bytes),
                        "nameSpaces": {},
                    },
                }
            ],
        },
        canonical=True,
    )


def _fake_ad() -> bytes:
    """A valid 4-element COSE_Sign1 stand-in, sized so the wrapped DeviceResponse
    exceeds 252 B and forces GET RESPONSE chaining (the real applet's 388 B
    case). The M2G.1 verifier doesn't crypto-verify the AD, only round-trips
    its bytes — so a 4-element CBOR array of the right size is enough."""
    return cbor2.dumps(
        [
            b"\xa0",            # protected (empty map, bstr-wrapped)
            {},                 # unprotected (empty map)
            b"\x11" * 200,      # payload — pad to push the response over CHUNK_LEN
            b"\x00" * 64,       # signature (raw r||s, fixed P-256 width)
        ],
        canonical=True,
    )


def _iv_reader(counter: int) -> bytes:
    return b"\x00" * 8 + counter.to_bytes(4, "big")


def _iv_device(counter: int) -> bytes:
    return b"\x00" * 7 + b"\x01" + counter.to_bytes(4, "big")


class _MockTransmit:
    """Records APDUs and replies from a scripted SELECT/EXCHANGE/ENVELOPE/
    GET RESPONSE sequence.

    EXCHANGE replies with the spec Reader Status sub-event RESPONSE
    ([0x01, 0x00, 0x00]) encrypted under SKDevice + deviceIv(1). The applet
    consumes deviceCounter=1 on this encrypt, so the subsequent ENVELOPE
    response uses deviceCounter=2.

    ENVELOPE replies with the first 252 B of a pre-encrypted DeviceResponse
    and SW=61xx; GET RESPONSE drains the rest."""

    CHUNK_LEN = 252
    EXCHANGE_RESPONSE_PT = b"\x01\x00\x00"

    def __init__(
        self,
        sk_device: bytes,
        sk_reader: bytes,
        device_response_plaintext: bytes,
    ):
        self.sk_device = sk_device
        self.sk_reader = sk_reader
        # Pre-encrypt the EXCHANGE response under SKDevice + deviceIv(1). The
        # applet's EXCHANGE handler consumes deviceCounter=1 first.
        self.exchange_ct = AESGCM(sk_device).encrypt(
            _iv_device(1), self.EXCHANGE_RESPONSE_PT, None
        )
        # Pre-encrypt the DeviceResponse under SKDevice + deviceIv(2) — after
        # EXCHANGE advanced the device counter from 1 -> 2.
        self.ct = AESGCM(sk_device).encrypt(
            _iv_device(2), device_response_plaintext, None
        )
        self.ct_off = 0
        self.apdus: list[bytes] = []
        self.reader_counter = 1
        self.device_counter = 1
        # The Reader Status sub-event REQUEST the verifier is expected to send.
        self.expected_exchange_pt = b"\x01\x00"
        # The DeviceRequest the verifier is expected to send.
        self.expected_request_pt = cbor2.dumps(
            {
                "version": "1.0",
                "docRequests": [
                    {
                        "itemsRequest": cbor2.dumps(
                            {
                                "docType": "org.iso.18013.5.1.mDL",
                                "nameSpaces": {},
                            },
                            canonical=True,
                        )
                    }
                ],
            },
            canonical=True,
        )

    def __call__(self, apdu: bytes) -> tuple[bytes, int]:
        self.apdus.append(apdu)
        cla, ins = apdu[0], apdu[1]

        # SELECT 5502 (CLA=0x00 INS=0xA4 P1=0x04 P2=0x00 Lc=09 ...)
        if cla == 0x00 and ins == 0xA4:
            return b"", 0x9000

        # EXCHANGE — decrypt [0x01, 0x00], reply with [0x01, 0x00, 0x00]
        # encrypted under SKDevice + deviceIv(device_counter=1).
        if cla == 0x80 and ins == 0xC9:
            lc = apdu[4]
            body = apdu[5 : 5 + lc]
            pt = AESGCM(self.sk_reader).decrypt(
                _iv_reader(self.reader_counter), body, None
            )
            assert pt == self.expected_exchange_pt, (
                f"EXCHANGE pt mismatch:\n  got={pt.hex()}\n  want={self.expected_exchange_pt.hex()}"
            )
            self.reader_counter += 1
            self.device_counter += 1
            return self.exchange_ct, 0x9000

        # ENVELOPE — decrypt DeviceRequest under reader-counter=2, ship first
        # chunk of the response encrypted under deviceCounter=2.
        if cla == 0x00 and ins == 0xC3:
            lc = apdu[4]
            body = apdu[5 : 5 + lc]
            pt = AESGCM(self.sk_reader).decrypt(
                _iv_reader(self.reader_counter), body, None
            )
            assert pt == self.expected_request_pt, (
                f"ENVELOPE pt mismatch:\n  got={pt.hex()}\n  want={self.expected_request_pt.hex()}"
            )
            self.reader_counter += 1
            chunk = self.ct[: self.CHUNK_LEN]
            self.ct_off = self.CHUNK_LEN
            remaining = len(self.ct) - self.CHUNK_LEN
            assert remaining > 0, "test fixture must produce a chained response"
            sw2 = 0xFF if remaining >= 0x100 else remaining
            return chunk, 0x6100 | sw2

        # GET RESPONSE — drain whatever's left.
        if cla == 0x00 and ins == 0xC0:
            chunk = self.ct[self.ct_off :]
            self.ct_off = len(self.ct)
            return chunk, 0x9000

        raise AssertionError(f"unexpected APDU: {apdu.hex()}")


def _setup_fixture() -> tuple[bytes, bytes, _MockTransmit]:
    step_up_sk = bytes(range(0xC0, 0xE0))
    sk_device, sk_reader = derive_step_up_session_keys(step_up_sk)
    ad = _fake_ad()
    device_response = _build_fixed_device_response(ad)
    mock = _MockTransmit(sk_device, sk_reader, device_response)
    return step_up_sk, ad, mock


def test_verify_step_up_m2_round_trips_access_document():
    step_up_sk, expected_ad, mock = _setup_fixture()

    ok, msg = verify_step_up_m2(mock, step_up_sk, expected_ad)

    assert ok, f"verify failed: {msg}"

    # Quick shape audit on the captured APDUs.
    sent = mock.apdus
    # SELECT, EXCHANGE, ENVELOPE, GET RESPONSE -> 4 APDUs.
    assert len(sent) == 4, [a.hex() for a in sent]
    assert sent[0][0:2] == bytes([0x00, 0xA4])   # SELECT
    assert sent[1][0:2] == bytes([0x80, 0xC9])   # EXCHANGE
    assert sent[2][0:2] == bytes([0x00, 0xC3])   # ENVELOPE
    assert sent[3][0:2] == bytes([0x00, 0xC0])   # GET RESPONSE
    # GET RESPONSE Le must equal the SW2 the applet handed us (136 remaining).
    expected_le = len(mock.ct) - _MockTransmit.CHUNK_LEN
    assert sent[3][4] == expected_le


def test_verify_step_up_m2_detects_ad_mismatch():
    step_up_sk, _, mock = _setup_fixture()
    # Pass a different AD than the one wrapped into the canned DeviceResponse.
    bogus_ad = cbor2.dumps([b"\xa0", {}, b"different", b"\x00" * 64], canonical=True)

    ok, msg = verify_step_up_m2(mock, step_up_sk, bogus_ad)

    assert not ok
    assert "issuerAuth" in msg or "Access Document" in msg or "AD" in msg, msg
