"""Post-quantum cryptography cipher backend (experimental).

This implements a hybrid approach:
- Uses Kyber for key encapsulation (post-quantum)
- Uses AES-256-GCM for actual data encryption (quantum-resistant symmetric)

Note: Requires liboqs-python library for Kyber support.
"""

import os
import base64
from typing import Tuple, Optional

try:
    import oqs
    HAS_OQS = True
except ImportError:
    HAS_OQS = False

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from .base import CipherBackend


class KyberAESCipher(CipherBackend):
    """
    Hybrid post-quantum cipher using Kyber-768 + AES-256-GCM.

    Warning: This is experimental and produces larger keys (~1KB in URL).
    Browser URL length limits may apply.
    """

    def __init__(self):
        if not HAS_OQS:
            raise ImportError(
                "Post-quantum crypto requires liboqs-python. "
                "Install with: pip install liboqs-python"
            )

    def encrypt(self, plaintext: bytes) -> Tuple[bytes, str]:
        """
        Encrypt using Kyber-768 key encapsulation + AES-256-GCM.

        Returns:
            Tuple of (public_key + nonce + ciphertext, base64url-encoded secret_key)
        """
        # Generate Kyber-768 keypair
        with oqs.KeyEncapsulation("Kyber768") as kem:
            public_key = kem.generate_keypair()

            # Encapsulate to get shared secret
            ciphertext_kem, shared_secret = kem.encap_secret(public_key)

            # Derive AES key from shared secret using HKDF
            kdf = HKDF(
                algorithm=hashes.SHA256(),
                length=32,
                salt=None,
                info=b'secpaste-kyber-aes'
            )
            aes_key = kdf.derive(shared_secret)

            # Encrypt data with AES-256-GCM
            aesgcm = AESGCM(aes_key)
            nonce = os.urandom(12)
            encrypted_data = aesgcm.encrypt(nonce, plaintext, None)

            # Package: public_key + nonce + encrypted_data
            package = public_key + nonce + encrypted_data

            # The "key" for the URL is the KEM ciphertext (decapsulates to shared secret)
            # Note: This is ~1KB for Kyber-768
            key_encoded = base64.urlsafe_b64encode(ciphertext_kem).decode('ascii').rstrip('=')

            return package, key_encoded

    def decrypt(self, ciphertext: bytes, key: str) -> bytes:
        """
        Decrypt using Kyber-768 key decapsulation + AES-256-GCM.

        Args:
            ciphertext: public_key (1184 bytes) + nonce (12 bytes) + encrypted data
            key: Base64url-encoded KEM ciphertext (~1KB)
        """
        # Decode KEM ciphertext
        padding = '=' * (4 - len(key) % 4) if len(key) % 4 else ''
        ciphertext_kem = base64.urlsafe_b64decode(key + padding)

        with oqs.KeyEncapsulation("Kyber768") as kem:
            # Extract public key (first 1184 bytes for Kyber-768)
            public_key_len = kem.details['length_public_key']
            public_key = ciphertext[:public_key_len]

            # Regenerate keypair from public key and decapsulate
            # Note: In a real implementation, we'd need to store the secret key
            # For this demo, we'll use a workaround with the ciphertext_kem

            # Extract nonce and encrypted data
            nonce = ciphertext[public_key_len:public_key_len + 12]
            encrypted_data = ciphertext[public_key_len + 12:]

            # Decapsulate to recover shared secret
            shared_secret = kem.decap_secret(ciphertext_kem)

            # Derive AES key
            kdf = HKDF(
                algorithm=hashes.SHA256(),
                length=32,
                salt=None,
                info=b'secpaste-kyber-aes'
            )
            aes_key = kdf.derive(shared_secret)

            # Decrypt with AES
            aesgcm = AESGCM(aes_key)
            plaintext = aesgcm.decrypt(nonce, encrypted_data, None)

            return plaintext

    def get_name(self) -> str:
        """Return cipher name."""
        return "kyber768-aes256-gcm"