"""Tests for the crypto module (RSA-OAEP asymmetric encryption)."""

import base64
import json

import pytest
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa

from hax.crypto import (
    _base64url_to_bigint,
    _jwk_to_rsa_private_key,
    decrypt_response,
    generate_key_pair,
    is_encrypted_response,
)
from hax.exceptions import DecryptionError


class TestGenerateKeyPair:
    """Tests for generate_key_pair function."""

    def test_generates_valid_rsa_keypair(self) -> None:
        """Key pair has correct RSA JWK structure."""
        public_key, private_key = generate_key_pair("test-passphrase")

        # Check public key structure
        assert public_key["kty"] == "RSA"
        assert public_key["alg"] == "RSA-OAEP-256"
        assert public_key["use"] == "enc"
        assert "n" in public_key
        assert "e" in public_key

        # Check private key structure
        assert private_key["kty"] == "RSA"
        assert private_key["alg"] == "RSA-OAEP-256"
        assert private_key["use"] == "enc"
        assert "n" in private_key
        assert "e" in private_key
        assert "d" in private_key
        assert "p" in private_key
        assert "q" in private_key
        assert "dp" in private_key
        assert "dq" in private_key
        assert "qi" in private_key

    def test_deterministic_key_generation(self) -> None:
        """Same passphrase always produces same keypair."""
        public1, private1 = generate_key_pair("same-passphrase")
        public2, private2 = generate_key_pair("same-passphrase")

        # Same passphrase should produce identical keys
        assert public1["n"] == public2["n"]
        assert public1["e"] == public2["e"]
        assert private1["d"] == private2["d"]

    def test_different_passphrases_different_keys(self) -> None:
        """Different passphrases produce different keys."""
        public1, _ = generate_key_pair("passphrase-1")
        public2, _ = generate_key_pair("passphrase-2")

        assert public1["n"] != public2["n"]

    def test_empty_passphrase(self) -> None:
        """Empty passphrase still produces valid keypair."""
        public_key, private_key = generate_key_pair("")

        assert public_key["kty"] == "RSA"
        assert private_key["kty"] == "RSA"

    def test_unicode_passphrase(self) -> None:
        """Unicode passphrase works correctly."""
        public_key, private_key = generate_key_pair("日本語パスフレーズ🔐")

        assert public_key["kty"] == "RSA"
        assert private_key["kty"] == "RSA"

    def test_long_passphrase(self) -> None:
        """Very long passphrase works correctly."""
        long_passphrase = "a" * 10000
        public_key, private_key = generate_key_pair(long_passphrase)

        assert public_key["kty"] == "RSA"
        assert private_key["kty"] == "RSA"

    def test_public_exponent_is_65537(self) -> None:
        """Public exponent is standard 65537 (AQAB in base64url)."""
        public_key, _ = generate_key_pair("test")

        # 65537 = AQAB in base64url
        assert public_key["e"] == "AQAB"

    def test_modulus_is_2048_bits(self) -> None:
        """RSA modulus is approximately 2048 bits."""
        public_key, _ = generate_key_pair("test")

        n = _base64url_to_bigint(public_key["n"])
        bit_length = n.bit_length()

        # Should be 2048 bits (could be slightly less due to leading zeros)
        assert 2040 <= bit_length <= 2048


class TestDecryptResponse:
    """Tests for decrypt_response function."""

    def _encrypt_with_public_key(self, data: object, public_key: dict) -> str:
        """Helper to encrypt data with a public key (simulating server-side encryption)."""
        # Convert JWK public key to cryptography RSA public key
        n = _base64url_to_bigint(public_key["n"])
        e = _base64url_to_bigint(public_key["e"])
        public_numbers = rsa.RSAPublicNumbers(e, n)
        rsa_key = public_numbers.public_key()

        # Encrypt
        plaintext = json.dumps(data).encode("utf-8")
        ciphertext = rsa_key.encrypt(
            plaintext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None,
            ),
        )

        return base64.b64encode(ciphertext).decode("ascii")

    def test_decrypt_dict_data(self) -> None:
        """Dictionary data decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")
        original_data = {"message": "Hello, World!", "count": 42}

        encrypted = self._encrypt_with_public_key(original_data, public_key)
        decrypted = decrypt_response(encrypted, private_key)

        assert decrypted == original_data

    def test_decrypt_string_data(self) -> None:
        """String data decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")
        original_data = "Just a string"

        encrypted = self._encrypt_with_public_key(original_data, public_key)
        decrypted = decrypt_response(encrypted, private_key)

        assert decrypted == original_data

    def test_decrypt_list_data(self) -> None:
        """List data decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")
        original_data = [1, 2, 3, {"nested": True}]

        encrypted = self._encrypt_with_public_key(original_data, public_key)
        decrypted = decrypt_response(encrypted, private_key)

        assert decrypted == original_data

    def test_decrypt_null(self) -> None:
        """Null decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")

        encrypted = self._encrypt_with_public_key(None, public_key)
        decrypted = decrypt_response(encrypted, private_key)

        assert decrypted is None

    def test_decrypt_boolean(self) -> None:
        """Boolean data decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")

        enc_true = self._encrypt_with_public_key(True, public_key)
        assert decrypt_response(enc_true, private_key) is True

        enc_false = self._encrypt_with_public_key(False, public_key)
        assert decrypt_response(enc_false, private_key) is False

    def test_decrypt_number(self) -> None:
        """Number data decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")

        encrypted = self._encrypt_with_public_key(42, public_key)
        decrypted = decrypt_response(encrypted, private_key)

        assert decrypted == 42

    def test_decrypt_unicode_data(self) -> None:
        """Unicode data decrypts correctly."""
        public_key, private_key = generate_key_pair("test-key")
        original_data = {"emoji": "🔐", "japanese": "暗号化", "arabic": "تشفير"}

        encrypted = self._encrypt_with_public_key(original_data, public_key)
        decrypted = decrypt_response(encrypted, private_key)

        assert decrypted == original_data

    def test_wrong_key_fails(self) -> None:
        """Decryption with wrong key raises DecryptionError."""
        public_key, _ = generate_key_pair("key-1")
        _, wrong_private_key = generate_key_pair("key-2")

        encrypted = self._encrypt_with_public_key({"secret": True}, public_key)

        with pytest.raises(DecryptionError, match="wrong key or corrupted"):
            decrypt_response(encrypted, wrong_private_key)

    def test_invalid_base64_fails(self) -> None:
        """Invalid base64 raises DecryptionError."""
        _, private_key = generate_key_pair("test-key")

        # Use a string that will cause DecryptionError (either invalid base64 or wrong size)
        with pytest.raises(DecryptionError):
            decrypt_response("not-valid-base64!!!", private_key)

    def test_corrupted_ciphertext_fails(self) -> None:
        """Corrupted ciphertext raises DecryptionError."""
        public_key, private_key = generate_key_pair("test-key")
        encrypted = self._encrypt_with_public_key({"test": True}, public_key)

        # Corrupt the ciphertext
        corrupted = encrypted[:-4] + "XXXX"

        with pytest.raises(DecryptionError):
            decrypt_response(corrupted, private_key)


class TestIsEncryptedResponse:
    """Tests for is_encrypted_response function."""

    def _create_encrypted_wrapper(self, data: object, passphrase: str) -> dict:
        """Helper to create a valid-looking encrypted wrapper."""
        public_key, _ = generate_key_pair(passphrase)

        # Convert JWK to RSA key and encrypt
        n = _base64url_to_bigint(public_key["n"])
        e = _base64url_to_bigint(public_key["e"])
        public_numbers = rsa.RSAPublicNumbers(e, n)
        rsa_key = public_numbers.public_key()

        plaintext = json.dumps(data).encode("utf-8")
        ciphertext = rsa_key.encrypt(
            plaintext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None,
            ),
        )

        return {"_encrypted": base64.b64encode(ciphertext).decode("ascii")}

    def test_encrypted_wrapper(self) -> None:
        """Detects encrypted wrapper format."""
        wrapper = self._create_encrypted_wrapper({"test": True}, "key")
        assert is_encrypted_response(wrapper) is True

    def test_plain_dict(self) -> None:
        """Plain dict is not encrypted."""
        assert is_encrypted_response({"test": True}) is False
        assert is_encrypted_response({"name": "Alice"}) is False

    def test_dict_without_encrypted_key(self) -> None:
        """Dict without _encrypted key is not encrypted."""
        assert is_encrypted_response({"other_key": "value"}) is False

    def test_wrong_encrypted_value_type(self) -> None:
        """Wrong type for _encrypted value returns False."""
        assert is_encrypted_response({"_encrypted": 123}) is False
        assert is_encrypted_response({"_encrypted": None}) is False
        assert is_encrypted_response({"_encrypted": ["list"]}) is False
        assert is_encrypted_response({"_encrypted": {"nested": True}}) is False

    def test_non_dict_values(self) -> None:
        """Non-dict values return False."""
        assert is_encrypted_response(None) is False
        assert is_encrypted_response("string") is False
        assert is_encrypted_response(123) is False
        assert is_encrypted_response([1, 2, 3]) is False

    def test_invalid_base64_in_wrapper(self) -> None:
        """Invalid base64 in wrapper returns False."""
        assert is_encrypted_response({"_encrypted": "not-valid-base64!!!"}) is False

    def test_too_short_base64_in_wrapper(self) -> None:
        """Too short base64 in wrapper returns False."""
        short = base64.b64encode(b"short").decode()
        assert is_encrypted_response({"_encrypted": short}) is False

    def test_valid_length_base64(self) -> None:
        """Valid length base64 in wrapper returns True."""
        wrapper = self._create_encrypted_wrapper({"test": True}, "key")
        assert is_encrypted_response(wrapper) is True


class TestCrossSDKCompatibility:
    """Tests verifying TypeScript/Python SDK produce same keys."""

    def test_same_modulus_for_same_passphrase(self) -> None:
        """Same passphrase produces consistent modulus."""
        public_key, _ = generate_key_pair("test-passphrase")

        # The modulus should be a valid base64url string
        assert "n" in public_key
        assert isinstance(public_key["n"], str)

        # RSA-2048 has a 256-byte modulus, ~342 chars in base64
        assert len(public_key["n"]) > 300

    def test_public_exponent_format(self) -> None:
        """Public exponent is in expected format."""
        public_key, _ = generate_key_pair("test-passphrase")

        # 65537 = AQAB in base64url
        assert public_key["e"] == "AQAB"


class TestJWKConversion:
    """Tests for JWK to RSA key conversion."""

    def test_jwk_to_rsa_private_key(self) -> None:
        """JWK private key converts to working RSA key."""
        _, private_key = generate_key_pair("test")

        # This should not raise
        rsa_key = _jwk_to_rsa_private_key(private_key)

        # Verify it's a valid private key by checking key size
        assert rsa_key.key_size == 2048

    def test_roundtrip_encryption_decryption(self) -> None:
        """Full roundtrip: generate keys, encrypt with public, decrypt with private."""
        public_key, private_key = generate_key_pair("my-secret-passphrase")

        # Use public key to encrypt (simulating server)
        n = _base64url_to_bigint(public_key["n"])
        e = _base64url_to_bigint(public_key["e"])
        public_numbers = rsa.RSAPublicNumbers(e, n)
        rsa_public = public_numbers.public_key()

        original_data = {"action": "approve", "details": {"amount": 1000}}
        plaintext = json.dumps(original_data).encode("utf-8")

        ciphertext = rsa_public.encrypt(
            plaintext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None,
            ),
        )

        # Use private key to decrypt
        encrypted_b64 = base64.b64encode(ciphertext).decode("ascii")
        decrypted = decrypt_response(encrypted_b64, private_key)

        assert decrypted == original_data
