"""Client-side encryption utilities for HAX SDK.

Uses RSA-OAEP for asymmetric encryption. The public key is sent to the server
with requests, and the server encrypts user responses with it. The private key
stays client-side and is used to decrypt responses.

Example:
    >>> from hax.crypto import generate_key_pair, decrypt_response
    >>>
    >>> # Generate keypair from passphrase (deterministic)
    >>> public_key, private_key = generate_key_pair("my-secret-passphrase")
    >>>
    >>> # Use public_key with HaxClient - it will be sent to server
    >>> # Server encrypts responses with this public key
    >>>
    >>> # Decrypt server response
    >>> plaintext = decrypt_response(encrypted_data, private_key)
"""

from __future__ import annotations

import base64
import hashlib
import json
from typing import Any, Callable

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

from hax.exceptions import DecryptionError

# RSA key size in bits
RSA_KEY_BITS = 2048

# Public exponent for RSA (65537 is standard)
RSA_PUBLIC_EXPONENT = 65537


def generate_key_pair(passphrase: str) -> tuple[dict[str, Any], dict[str, Any]]:
    """Generate a deterministic RSA keypair from a passphrase.

    The same passphrase will always produce the same keypair, allowing
    users to recover their keys from a remembered passphrase.

    Args:
        passphrase: String passphrase to derive keys from.

    Returns:
        Tuple of (public_key_jwk, private_key_jwk) in JWK format.

    Example:
        >>> public_key, private_key = generate_key_pair("my-secret")
        >>> public_key["kty"]
        'RSA'
    """
    # Create deterministic seed from passphrase
    seed = hashlib.sha256(passphrase.encode("utf-8")).digest()

    # Generate RSA key components deterministically
    # We use a seeded PRNG to generate prime candidates
    prng = _create_seeded_prng(seed)

    # Generate two large primes p and q
    bit_size = RSA_KEY_BITS // 2
    p = _generate_prime(bit_size, prng)
    q = _generate_prime(bit_size, prng)

    # Ensure p > q for standard RSA
    if p < q:
        p, q = q, p

    # Calculate RSA components
    n = p * q
    phi = (p - 1) * (q - 1)
    e = RSA_PUBLIC_EXPONENT
    d = _mod_inverse(e, phi)

    # Calculate CRT components for efficiency
    dp = d % (p - 1)
    dq = d % (q - 1)
    qi = _mod_inverse(q, p)

    # Convert to JWK format
    public_key: dict[str, Any] = {
        "kty": "RSA",
        "alg": "RSA-OAEP-256",
        "use": "enc",
        "n": _bigint_to_base64url(n),
        "e": _bigint_to_base64url(e),
    }

    private_key: dict[str, Any] = {
        "kty": "RSA",
        "alg": "RSA-OAEP-256",
        "use": "enc",
        "n": _bigint_to_base64url(n),
        "e": _bigint_to_base64url(e),
        "d": _bigint_to_base64url(d),
        "p": _bigint_to_base64url(p),
        "q": _bigint_to_base64url(q),
        "dp": _bigint_to_base64url(dp),
        "dq": _bigint_to_base64url(dq),
        "qi": _bigint_to_base64url(qi),
    }

    return public_key, private_key


def decrypt_response(encrypted: str, private_key: dict[str, Any]) -> Any:
    """Decrypt a response that was encrypted by the server using RSA-OAEP.

    The server encrypts responses using the public key provided when creating
    the request. This function decrypts using the corresponding private key.

    Args:
        encrypted: Base64-encoded encrypted data from server.
        private_key: JWK private key (must match the public key sent to server).

    Returns:
        Decrypted data (original JSON structure).

    Raises:
        DecryptionError: If decryption fails (wrong key, corrupted data).

    Example:
        >>> _, private_key = generate_key_pair("my-secret")
        >>> # encrypted = ... (from server)
        >>> data = decrypt_response(encrypted, private_key)
    """
    # Convert JWK to cryptography RSA private key
    rsa_key = _jwk_to_rsa_private_key(private_key)

    # Decode base64
    try:
        ciphertext = _base64_decode(encrypted)
    except Exception as e:
        raise DecryptionError(f"Invalid base64 encoding: {e}") from e

    # Decrypt
    try:
        plaintext = rsa_key.decrypt(
            ciphertext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None,
            ),
        )
    except Exception as e:
        raise DecryptionError(f"Decryption failed (wrong key or corrupted data): {e}") from e

    # Parse JSON
    try:
        return json.loads(plaintext.decode("utf-8"))
    except json.JSONDecodeError as e:
        raise DecryptionError(f"Invalid JSON in decrypted data: {e}") from e


def is_encrypted_response(data: Any) -> bool:
    """Check if data looks like it was encrypted by the server.

    Encrypted responses are stored as {"_encrypted": "base64..."} wrapper.

    Args:
        data: Data to check.

    Returns:
        True if data appears to be encrypted, False otherwise.

    Example:
        >>> is_encrypted_response({"_encrypted": "abc..."})
        True
        >>> is_encrypted_response({"test": True})
        False
    """
    if not isinstance(data, dict):
        return False
    if "_encrypted" not in data:
        return False
    enc_value = data.get("_encrypted")
    if not isinstance(enc_value, str):
        return False
    try:
        raw = _base64_decode(enc_value)
        # RSA-2048 produces 256 bytes of ciphertext
        return 128 <= len(raw) <= 512
    except Exception:
        return False


# ─────────────────────────────────────────────────────────────────────────────
# Internal Utilities
# ─────────────────────────────────────────────────────────────────────────────


def _create_seeded_prng(seed: bytes) -> Callable[[int], bytes]:
    """Create a seeded PRNG that produces deterministic random bytes."""
    state = seed
    counter = 0

    def get_random_bytes(length: int = 32) -> bytes:
        nonlocal state, counter
        chunks = []
        remaining = length

        while remaining > 0:
            # Hash state + counter to get next random bytes
            counter_bytes = counter.to_bytes(4, "little")
            hash_val = hashlib.sha256(state + counter_bytes).digest()

            take = min(remaining, len(hash_val))
            chunks.append(hash_val[:take])
            remaining -= take
            counter += 1

            # Update state periodically
            if counter % 256 == 0:
                state = hashlib.sha256(state + hash_val).digest()

        return b"".join(chunks)

    return get_random_bytes


def _generate_prime(bits: int, prng: Callable[[int], bytes]) -> int:
    """Generate a random prime of the specified bit size using Miller-Rabin."""
    byte_length = (bits + 7) // 8

    while True:
        random_bytes = bytearray(prng(byte_length))

        # Set MSB to ensure we get a number of the right bit length
        random_bytes[0] |= 0x80
        # Set LSB to ensure odd number
        random_bytes[-1] |= 0x01

        candidate = int.from_bytes(random_bytes, "big")

        # Quick check against small primes
        if not _passes_small_prime_check(candidate):
            continue

        # Miller-Rabin primality test with 40 rounds
        if _miller_rabin(candidate, 40, prng):
            return candidate


_SMALL_PRIMES = [
    3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53,
    59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113,
]


def _passes_small_prime_check(n: int) -> bool:
    """Check if n is divisible by small primes."""
    for p in _SMALL_PRIMES:
        if n == p:
            return True
        if n % p == 0:
            return False
    return True


def _miller_rabin(n: int, rounds: int, prng: Callable[[int], bytes]) -> bool:
    """Miller-Rabin primality test."""
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False

    # Write n-1 as 2^r * d
    r = 0
    d = n - 1
    while d % 2 == 0:
        r += 1
        d //= 2

    # Witness loop
    byte_length = (n.bit_length() + 7) // 8

    for _ in range(rounds):
        # Generate random witness a in range [2, n-2]
        while True:
            random_bytes = prng(byte_length)
            a = int.from_bytes(random_bytes, "big") % (n - 3) + 2
            if 2 <= a < n - 1:
                break

        x = pow(a, d, n)

        if x == 1 or x == n - 1:
            continue

        composite = True
        for _ in range(r - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                composite = False
                break

        if composite:
            return False

    return True


def _mod_inverse(a: int, m: int) -> int:
    """Extended Euclidean Algorithm to find modular inverse."""
    old_r, r = a, m
    old_s, s = 1, 0

    while r != 0:
        quotient = old_r // r
        old_r, r = r, old_r - quotient * r
        old_s, s = s, old_s - quotient * s

    if old_r > 1:
        raise ValueError("Modular inverse does not exist")

    return old_s % m


def _bigint_to_base64url(n: int) -> str:
    """Convert integer to base64url encoding (for JWK)."""
    if n == 0:
        return "AA"

    # Calculate byte length
    byte_length = (n.bit_length() + 7) // 8
    data = n.to_bytes(byte_length, "big")

    # Base64url encode (no padding)
    encoded = base64.urlsafe_b64encode(data).decode("ascii")
    return encoded.rstrip("=")


def _base64url_to_bigint(s: str) -> int:
    """Convert base64url encoding to integer."""
    # Add padding if needed
    padding_needed = 4 - (len(s) % 4)
    if padding_needed != 4:
        s += "=" * padding_needed

    data = base64.urlsafe_b64decode(s)
    return int.from_bytes(data, "big")


def _base64_decode(data: str) -> bytes:
    """Base64 decode (handles both standard and url-safe)."""
    # Convert base64url to standard base64
    data = data.replace("-", "+").replace("_", "/")
    # Add padding if needed
    padding_needed = 4 - (len(data) % 4)
    if padding_needed != 4:
        data += "=" * padding_needed

    return base64.b64decode(data)


def _jwk_to_rsa_private_key(jwk: dict[str, Any]) -> rsa.RSAPrivateKey:
    """Convert a JWK to a cryptography RSA private key."""
    # Extract components
    n = _base64url_to_bigint(jwk["n"])
    e = _base64url_to_bigint(jwk["e"])
    d = _base64url_to_bigint(jwk["d"])
    p = _base64url_to_bigint(jwk["p"])
    q = _base64url_to_bigint(jwk["q"])
    dp = _base64url_to_bigint(jwk["dp"])
    dq = _base64url_to_bigint(jwk["dq"])
    qi = _base64url_to_bigint(jwk["qi"])

    # Create RSA private key
    public_numbers = rsa.RSAPublicNumbers(e, n)
    private_numbers = rsa.RSAPrivateNumbers(p, q, d, dp, dq, qi, public_numbers)

    return private_numbers.private_key()
