"""Tests for pm3py.sim.crypto1 and pm3py.sim.mifare — Crypto-1 cipher and MIFARE Classic."""
import asyncio
import struct
import pytest

from pm3py.sim.frame import RFFrame
from pm3py.sim.medium import SoftwareMedium
from pm3py.transponders.hf.iso14443a.nxp.crypto1 import Crypto1
from pm3py.transponders.hf.iso14443a.nxp.mifare_classic import MifareClassicTag, MifareClassicReader
from pm3py.transponders.hf.iso14443a.base import Reader14443A, REQA, CL1, NVB_SELECT, _compute_bcc


def run(coro):
    return asyncio.get_event_loop().run_until_complete(coro)


# ---------------------------------------------------------------------------
# Crypto-1 LFSR
# ---------------------------------------------------------------------------

class TestCrypto1LFSR:
    """Test Crypto-1 cipher against known test vectors."""

    def test_key_load(self):
        """Verify LFSR state after loading a key."""
        c = Crypto1(b"\xFF\xFF\xFF\xFF\xFF\xFF")
        # After loading all-ones key, LFSR should be non-zero
        assert c._lfsr != 0

    def test_generate_keystream_is_deterministic(self):
        """Same key produces same keystream."""
        c1 = Crypto1(b"\xA0\xA1\xA2\xA3\xA4\xA5")
        c2 = Crypto1(b"\xA0\xA1\xA2\xA3\xA4\xA5")
        ks1 = [c1.generate_bit() for _ in range(32)]
        ks2 = [c2.generate_bit() for _ in range(32)]
        assert ks1 == ks2

    def test_different_keys_different_keystream(self):
        c1 = Crypto1(b"\x00\x00\x00\x00\x00\x00")
        c2 = Crypto1(b"\xFF\xFF\xFF\xFF\xFF\xFF")
        ks1 = [c1.generate_bit() for _ in range(32)]
        ks2 = [c2.generate_bit() for _ in range(32)]
        assert ks1 != ks2

    def test_encrypt_decrypt_roundtrip(self):
        """Encrypting then decrypting with same state should return original."""
        key = b"\xA0\xA1\xA2\xA3\xA4\xA5"
        plaintext = b"\xDE\xAD\xBE\xEF"

        c_enc = Crypto1(key)
        encrypted = c_enc.encrypt_bytes(plaintext)

        c_dec = Crypto1(key)
        decrypted = c_dec.encrypt_bytes(encrypted)  # XOR cipher: enc == dec

        assert decrypted == plaintext

    def test_known_vector_lfsr_feedback(self):
        """Verify the LFSR feedback polynomial is correct.

        The Crypto-1 LFSR polynomial is:
        x^48 + x^43 + x^39 + x^38 + x^36 + x^34 + x^33 + x^31 + x^29 +
        x^24 + x^23 + x^21 + x^19 + x^13 + x^9 + x^7 + x^6 + x^5 + 1
        """
        c = Crypto1(b"\x00\x00\x00\x00\x00\x01")
        # Just verify it doesn't crash and produces deterministic output
        bits = [c.generate_bit() for _ in range(48)]
        assert len(bits) == 48


class TestCrypto1Auth:
    """Test Crypto-1 mutual authentication protocol."""

    def test_tag_nonce_generation(self):
        """Tag generates a 4-byte nonce."""
        c = Crypto1(b"\xFF\xFF\xFF\xFF\xFF\xFF")
        nt = c.generate_nonce()
        assert len(nt) == 4

    def test_auth_mutual_success(self):
        """Full mutual authentication between tag and reader Crypto-1 instances."""
        key = b"\xFF\xFF\xFF\xFF\xFF\xFF"
        uid = b"\x01\x02\x03\x04"

        # Tag side: generate nonce
        tag_crypto = Crypto1(key)
        nt = tag_crypto.generate_nonce()

        # Initialize both sides with uid ^ nt
        uid_int = struct.unpack(">I", uid)[0]
        nt_int = struct.unpack(">I", nt)[0]

        tag_crypto = Crypto1(key)
        tag_crypto.init_auth(uid_int, nt_int)

        reader_crypto = Crypto1(key)
        reader_crypto.init_auth(uid_int, nt_int)

        # Reader generates nr (random) and computes ar = suc(nt, 64)
        nr = b"\xAB\xCD\xEF\x01"

        # Encrypt nr with reader's keystream
        nr_enc = reader_crypto.encrypt_bytes(nr)
        # Reader computes ar (encrypted successor of nt)
        ar = reader_crypto.encrypt_bytes(struct.pack(">I", _suc(nt_int, 64)))

        # Tag decrypts nr
        nr_dec = tag_crypto.encrypt_bytes(nr_enc)
        assert nr_dec == nr

        # Tag decrypts ar and verifies
        ar_dec_int = struct.unpack(">I", tag_crypto.encrypt_bytes(ar))[0]
        assert ar_dec_int == _suc(nt_int, 64)

        # Tag sends at (encrypted successor of nt, 96)
        at = tag_crypto.encrypt_bytes(struct.pack(">I", _suc(nt_int, 96)))

        # Reader verifies at
        at_dec_int = struct.unpack(">I", reader_crypto.encrypt_bytes(at))[0]
        assert at_dec_int == _suc(nt_int, 96)


def _suc(nt: int, n: int) -> int:
    """Compute successor of nt by n LFSR clocks (simplified for test)."""
    # In real Crypto-1, suc is the LFSR state after n clocks
    # For testing, we use a simple PRNG-like computation
    val = nt
    for _ in range(n):
        bit = ((val >> 31) ^ (val >> 20) ^ (val >> 15) ^ (val >> 0)) & 1
        val = ((val << 1) | bit) & 0xFFFFFFFF
    return val


# ---------------------------------------------------------------------------
# MifareClassicTag — transponder model
# ---------------------------------------------------------------------------

class TestMifareClassicTagBasics:
    """Test MifareClassicTag construction and 14443-A compliance."""

    def test_1k_atqa_sak(self):
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        run(tag.power_on())
        resp = run(tag.handle_frame(RFFrame.from_hex("26")))
        assert resp.data == b"\x04\x00"  # ATQA for 1K

        # SELECT
        uid = b"\x01\x02\x03\x04"
        bcc = _compute_bcc(uid)
        resp = run(tag.handle_frame(RFFrame.from_bytes(b"\x93\x70" + uid + bytes([bcc]))))
        assert resp.data[0] == 0x08  # SAK for 1K

    def test_4k_atqa_sak(self):
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="4k")
        run(tag.power_on())
        resp = run(tag.handle_frame(RFFrame.from_hex("26")))
        assert resp.data == b"\x02\x00"  # ATQA for 4K

        uid = b"\x01\x02\x03\x04"
        bcc = _compute_bcc(uid)
        resp = run(tag.handle_frame(RFFrame.from_bytes(b"\x93\x70" + uid + bytes([bcc]))))
        assert resp.data[0] == 0x18  # SAK for 4K

    def test_rejects_rats(self):
        """MIFARE Classic is Part 3 only."""
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        run(tag.power_on())
        run(tag.handle_frame(RFFrame.from_hex("26")))
        uid = b"\x01\x02\x03\x04"
        bcc = _compute_bcc(uid)
        run(tag.handle_frame(RFFrame.from_bytes(b"\x93\x70" + uid + bytes([bcc]))))

        resp = run(tag.handle_frame(RFFrame.from_hex("E050")))
        assert resp is None

    def test_default_memory_all_zeros(self):
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        assert len(tag._data) == 1024  # 64 blocks * 16 bytes
        assert tag._data[:16] == b"\x00" * 16

    def test_4k_memory_size(self):
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="4k")
        assert len(tag._data) == 4096  # 256 blocks * 16 bytes

    def test_default_keys_are_ff(self):
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04")
        for sector in range(16):
            assert tag._keys_a[sector] == b"\xFF" * 6
            assert tag._keys_b[sector] == b"\xFF" * 6


class TestMifareClassicTagAuth:
    """Test MIFARE Classic authentication."""

    def test_auth_command_returns_tag_nonce(self):
        """AUTH command (0x60/0x61) should return a 4-byte encrypted tag nonce."""
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        run(tag.power_on())
        run(tag.handle_frame(RFFrame.from_hex("26")))
        uid = b"\x01\x02\x03\x04"
        bcc = _compute_bcc(uid)
        run(tag.handle_frame(RFFrame.from_bytes(b"\x93\x70" + uid + bytes([bcc]))))

        # AUTH_A for block 0
        resp = run(tag.handle_frame(RFFrame.from_bytes(b"\x60\x00")))
        assert resp is not None
        assert len(resp.data) == 4  # tag nonce (nt)

    def test_unauthenticated_read_rejected(self):
        """READ without authentication should be rejected."""
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        run(tag.power_on())
        run(tag.handle_frame(RFFrame.from_hex("26")))
        uid = b"\x01\x02\x03\x04"
        bcc = _compute_bcc(uid)
        run(tag.handle_frame(RFFrame.from_bytes(b"\x93\x70" + uid + bytes([bcc]))))

        # READ block 0 without auth
        resp = run(tag.handle_frame(RFFrame.from_bytes(b"\x30\x00")))
        assert resp is None  # rejected


class TestMifareClassicTagMemory:
    """Test MIFARE Classic memory layout."""

    def test_1k_sector_block_mapping(self):
        """1K: sectors 0-15, 4 blocks each."""
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        assert tag.sector_for_block(0) == 0
        assert tag.sector_for_block(3) == 0
        assert tag.sector_for_block(4) == 1
        assert tag.sector_for_block(63) == 15

    def test_4k_sector_block_mapping(self):
        """4K: sectors 0-31 = 4 blocks, sectors 32-39 = 16 blocks."""
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="4k")
        assert tag.sector_for_block(0) == 0
        assert tag.sector_for_block(127) == 31  # last block of sector 31
        assert tag.sector_for_block(128) == 32  # first block of sector 32
        assert tag.sector_for_block(255) == 39  # last block of sector 39

    def test_read_block_data(self):
        """Direct read of block data from memory."""
        data = bytearray(1024)
        data[0:16] = b"\xDE\xAD\xBE\xEF" + b"\x00" * 12
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k", data=data)
        assert tag.read_block_raw(0) == b"\xDE\xAD\xBE\xEF" + b"\x00" * 12

    def test_write_block_data(self):
        """Direct write of block data to memory."""
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        new_data = b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\x0F\x10"
        tag.write_block_raw(4, new_data)
        assert tag.read_block_raw(4) == new_data


class TestMifareClassicReader:
    """Test MifareClassicReader with full auth + read/write."""

    def test_read_block_with_default_key(self):
        """Reader authenticates and reads a block."""
        medium = SoftwareMedium()
        data = bytearray(1024)
        data[16:32] = b"\xCA\xFE\xBA\xBE" + b"\x00" * 12
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k", data=data)
        run(medium.attach(tag))

        reader = MifareClassicReader(medium)
        result = run(reader.read_block(
            uid=b"\x01\x02\x03\x04",
            block=1,
            key=b"\xFF\xFF\xFF\xFF\xFF\xFF",
        ))
        assert result["success"]
        assert result["data"][:4] == b"\xCA\xFE\xBA\xBE"

    def test_write_block_with_default_key(self):
        """Reader authenticates and writes a block."""
        medium = SoftwareMedium()
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        run(medium.attach(tag))

        reader = MifareClassicReader(medium)
        new_data = b"\x01" * 16
        result = run(reader.write_block(
            uid=b"\x01\x02\x03\x04",
            block=4,
            data=new_data,
            key=b"\xFF\xFF\xFF\xFF\xFF\xFF",
        ))
        assert result["success"]

        # Verify by reading back
        result = run(reader.read_block(
            uid=b"\x01\x02\x03\x04",
            block=4,
            key=b"\xFF\xFF\xFF\xFF\xFF\xFF",
        ))
        assert result["data"] == new_data

    def test_wrong_key_fails(self):
        """Reader with wrong key cannot authenticate."""
        medium = SoftwareMedium()
        tag = MifareClassicTag(uid=b"\x01\x02\x03\x04", size="1k")
        run(medium.attach(tag))

        reader = MifareClassicReader(medium)
        result = run(reader.read_block(
            uid=b"\x01\x02\x03\x04",
            block=0,
            key=b"\x00\x00\x00\x00\x00\x00",  # wrong key
        ))
        assert not result["success"]