"""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.sim.crypto1 import Crypto1
from pm3py.sim.mifare import MifareClassicTag, MifareClassicReader
from pm3py.sim.iso14443a 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"]