"""
Pure cryptographic primitives for Kilit.

This module contains NO IO operations, NO environment variable reading,
and NO dependencies on higher layers. It provides pure cryptographic
functions for encryption, decryption, and key generation.
"""

from __future__ import annotations

import base64
import os
from dataclasses import dataclass
from typing import Any

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt


@dataclass
class Envelope:
    """
    Represents the encrypted file envelope format (version 1).

    The envelope contains all metadata needed for decryption:
    - version: Format version (currently 1)
    - kdf: Key derivation function name (currently "scrypt")
    - kdfp: KDF parameters dict with n, r, p (CRITICAL: ensures deterministic decryption)
    - salt: Random salt for Scrypt (base64 urlsafe encoded)
    - nonce: Random nonce for AES-GCM (base64 urlsafe encoded)
    - ciphertext: Encrypted data (base64 urlsafe encoded)

    Versioning Strategy:
    - Version 1: Current format with kdfp field
    - Future versions can add fields while maintaining backward compatibility
    - Decryption validates version and rejects unsupported versions

    Security Note:
    - Storing kdfp in envelope ensures files encrypted with different parameters
      remain decryptable regardless of environment variable changes
    - This prevents "environment drift" from breaking decryption in CI/prod/local
    """

    version: int
    kdf: str
    kdfp: dict[str, int]  # {"n": int, "r": int, "p": int}
    salt: str  # base64 urlsafe encoded
    nonce: str  # base64 urlsafe encoded
    ciphertext: str  # base64 urlsafe encoded

    def to_dict(self) -> dict[str, Any]:
        """
        Convert envelope to dictionary for JSON serialization.

        Uses short field names to minimize file size:
        - v: version
        - kdf: key derivation function
        - kdfp: KDF parameters
        - salt: salt
        - nonce: nonce
        - ct: ciphertext
        """
        return {
            "v": self.version,
            "kdf": self.kdf,
            "kdfp": self.kdfp,
            "salt": self.salt,
            "nonce": self.nonce,
            "ct": self.ciphertext,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Envelope:
        """
        Create envelope from dictionary with strict validation.

        Args:
            data: Dictionary with envelope fields

        Returns:
            Envelope instance

        Raises:
            ValueError: If version is not 1 or kdfp structure is invalid
            KeyError: If required fields are missing

        Validation:
        - Checks all required fields present (v, kdf, kdfp, salt, nonce, ct)
        - Validates version == 1
        - Validates kdfp contains n, r, p as integers
        - Validates base64 format (urlsafe encoding)
        """
        # Check required fields
        required_fields = ["v", "kdf", "kdfp", "salt", "nonce", "ct"]
        missing = [f for f in required_fields if f not in data]
        if missing:
            raise KeyError(f"Missing required envelope fields: {', '.join(missing)}")

        # Validate version
        version = data["v"]
        if version != 1:
            raise ValueError(
                f"Unsupported envelope version {version}. This package supports version 1 only."
            )

        # Validate kdfp structure
        kdfp = data["kdfp"]
        if not isinstance(kdfp, dict):
            raise ValueError(f"Invalid kdfp: expected dict, got {type(kdfp).__name__}")

        required_kdfp_fields = ["n", "r", "p"]
        missing_kdfp = [f for f in required_kdfp_fields if f not in kdfp]
        if missing_kdfp:
            raise ValueError(f"Missing kdfp fields: {', '.join(missing_kdfp)}")

        for field in required_kdfp_fields:
            if not isinstance(kdfp[field], int):
                raise ValueError(
                    f"Invalid kdfp.{field}: expected int, got {type(kdfp[field]).__name__}"
                )

        # Validate base64 format (urlsafe)
        for field_name, field_value in [
            ("salt", data["salt"]),
            ("nonce", data["nonce"]),
            ("ct", data["ct"]),
        ]:
            try:
                # Attempt to decode to validate format (add padding if needed)
                padding = "=" * (4 - len(field_value) % 4) if len(field_value) % 4 else ""
                base64.urlsafe_b64decode(field_value + padding)
            except Exception as e:
                raise ValueError(f"Invalid base64 encoding in field '{field_name}': {e}")

        return cls(
            version=version,
            kdf=data["kdf"],
            kdfp=kdfp,
            salt=data["salt"],
            nonce=data["nonce"],
            ciphertext=data["ct"],
        )


def validate_scrypt_params(n: int, r: int, p: int) -> None:
    """
    Validate Scrypt parameters for security and correctness.

    Args:
        n: CPU/memory cost parameter (must be power of 2, >= 2**14)
        r: Block size parameter (must be >= 1)
        p: Parallelization parameter (must be >= 1)

    Raises:
        ValueError: If any parameter is invalid

    Security Assumptions:
    - n >= 2**14 (16384): Minimum secure threshold for modern hardware
    - n must be power of 2: Required by Scrypt algorithm
    - r >= 1, p >= 1: Algorithm requirements

    Rationale for n >= 2**14:
    - Provides reasonable protection against brute-force attacks
    - Balances security with performance for typical use cases
    - Can be increased via environment variables for higher security needs

    References:
    - RFC 7914: The scrypt Password-Based Key Derivation Function
    - OWASP recommendations for password hashing
    """
    # Validate n is power of 2
    if n <= 0 or (n & (n - 1)) != 0:
        raise ValueError(
            f"Scrypt parameter 'n' must be a power of 2, got {n}. "
            f"Valid values: 2**14 (16384), 2**15 (32768), 2**16 (65536), etc."
        )

    # Validate n meets minimum security threshold
    min_n = 2**14  # 16384
    if n < min_n:
        raise ValueError(
            f"Scrypt parameter 'n' must be >= {min_n} for security, got {n}. "
            f"Lower values provide insufficient protection against brute-force attacks."
        )

    # Validate r >= 1
    if r < 1:
        raise ValueError(f"Scrypt parameter 'r' must be >= 1, got {r}")

    # Validate p >= 1
    if p < 1:
        raise ValueError(f"Scrypt parameter 'p' must be >= 1, got {p}")


class Encryptor:
    """
    Pure cryptographic encryptor using AES-256-GCM with Scrypt KDF.

    This class is PURE - it performs NO IO operations and reads NO environment variables.
    All parameters must be provided explicitly by the caller (service layer).

    Encryption:
    - Generates random salt and nonce for each operation
    - Derives encryption key using Scrypt(master_key, salt, n, r, p)
    - Encrypts using AES-256-GCM with AAD=b"kilit-v1"
    - Stores Scrypt parameters in envelope (kdfp field)
    - Returns envelope dict with all metadata

    Decryption:
    - Extracts Scrypt parameters from envelope.kdfp (NOT from environment)
    - Derives encryption key using envelope salt and envelope.kdfp
    - Decrypts using AES-256-GCM with AAD=b"kilit-v1"
    - Returns plaintext string

    Security Properties:
    - Non-deterministic encryption (random salt/nonce per operation)
    - Authenticated encryption (AES-GCM prevents tampering)
    - Forward secrecy (each file has unique salt)
    - Deterministic decryption (uses envelope.kdfp, not environment)

    Thread Safety:
    - This class is stateless and thread-safe
    - Each method call is independent
    """

    # AAD (Additional Authenticated Data) for AES-GCM
    # This prevents encrypted files from being used in different contexts
    AAD = b"kilit-v1"

    def __init__(self, master_key: bytes, n: int = 2**14, r: int = 8, p: int = 1):
        """
        Initialize encryptor with master key and Scrypt parameters.

        Args:
            master_key: 32-byte master key (decoded from base64)
            n: Scrypt CPU/memory cost (default: 2**14 = 16384)
            r: Scrypt block size (default: 8)
            p: Scrypt parallelization (default: 1)

        Raises:
            ValueError: If Scrypt parameters are invalid

        Note:
            These parameters are used for NEW encryptions only.
            Decryption always uses parameters from the envelope.
        """
        validate_scrypt_params(n, r, p)
        self.master_key = master_key
        self.n = n
        self.r = r
        self.p = p

    def encrypt(self, plaintext: str) -> dict[str, Any]:
        """
        Encrypt plaintext and return envelope dictionary.

        Args:
            plaintext: String to encrypt

        Returns:
            Envelope dictionary with fields: v, kdf, kdfp, salt, nonce, ct

        Process:
        1. Generate random 16-byte salt
        2. Generate random 12-byte nonce
        3. Derive 32-byte key using Scrypt(master_key, salt, n, r, p)
        4. Encrypt using AES-256-GCM with AAD
        5. Encode all binary data as base64 urlsafe (no padding)
        6. Store Scrypt parameters in envelope.kdfp
        7. Return envelope dict

        Security:
        - Each encryption produces different ciphertext (random salt/nonce)
        - Scrypt parameters stored in envelope ensure deterministic decryption
        - AAD prevents ciphertext from being used in different contexts
        """
        # Generate random salt and nonce
        salt = os.urandom(16)
        nonce = os.urandom(12)

        # Derive encryption key using Scrypt
        kdf = Scrypt(salt=salt, length=32, n=self.n, r=self.r, p=self.p)
        key = kdf.derive(self.master_key)

        # Encrypt using AES-256-GCM
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, plaintext.encode("utf-8"), self.AAD)

        # Encode as base64 urlsafe (no padding)
        salt_b64 = base64.urlsafe_b64encode(salt).decode("ascii").rstrip("=")
        nonce_b64 = base64.urlsafe_b64encode(nonce).decode("ascii").rstrip("=")
        ct_b64 = base64.urlsafe_b64encode(ciphertext).decode("ascii").rstrip("=")

        # Create envelope with Scrypt parameters
        envelope = Envelope(
            version=1,
            kdf="scrypt",
            kdfp={"n": self.n, "r": self.r, "p": self.p},
            salt=salt_b64,
            nonce=nonce_b64,
            ciphertext=ct_b64,
        )

        return envelope.to_dict()

    def decrypt(self, envelope_dict: dict[str, Any]) -> str:
        """
        Decrypt envelope and return plaintext.

        Args:
            envelope_dict: Envelope dictionary from encrypted file

        Returns:
            Decrypted plaintext string

        Raises:
            ValueError: If envelope version is unsupported or structure is invalid
            InvalidTag: If authentication fails (wrong key or corrupted data)

        Process:
        1. Parse and validate envelope
        2. Extract Scrypt parameters from envelope.kdfp (NOT from environment)
        3. Validate extracted parameters
        4. Decode base64 fields (add padding if needed)
        5. Derive key using envelope salt and envelope.kdfp
        6. Decrypt using AES-256-GCM with AAD
        7. Return plaintext

        Security:
        - Uses envelope.kdfp for deterministic decryption
        - Environment variables do NOT affect decryption
        - This ensures files remain decryptable regardless of environment changes
        - InvalidTag exception indicates wrong key or corrupted data
        """
        # Parse and validate envelope
        envelope = Envelope.from_dict(envelope_dict)

        # Extract and validate Scrypt parameters from envelope
        n = envelope.kdfp["n"]
        r = envelope.kdfp["r"]
        p = envelope.kdfp["p"]
        validate_scrypt_params(n, r, p)

        # Decode base64 fields (add padding if needed)
        def decode_b64(s: str) -> bytes:
            # Add padding if needed (base64 requires length % 4 == 0)
            padding = "=" * (4 - len(s) % 4) if len(s) % 4 else ""
            return base64.urlsafe_b64decode(s + padding)

        salt = decode_b64(envelope.salt)
        nonce = decode_b64(envelope.nonce)
        ciphertext = decode_b64(envelope.ciphertext)

        # Derive encryption key using envelope salt and envelope.kdfp
        kdf = Scrypt(salt=salt, length=32, n=n, r=r, p=p)
        key = kdf.derive(self.master_key)

        # Decrypt using AES-256-GCM
        aesgcm = AESGCM(key)
        plaintext_bytes = aesgcm.decrypt(nonce, ciphertext, self.AAD)

        return plaintext_bytes.decode("utf-8")

    def __repr__(self) -> str:
        """
        Safe repr that doesn't expose master key.

        Returns:
            String representation without sensitive data
        """
        return f"Encryptor(n={self.n}, r={self.r}, p={self.p})"


def generate_master_key() -> str:
    """
    Generate a new 32-byte master key.

    Returns:
        Base64 urlsafe encoded master key (no padding)

    Process:
    1. Generate 32 random bytes using os.urandom (cryptographically secure)
    2. Encode as base64 urlsafe
    3. Remove padding (=) for cleaner format
    4. Return as string

    Format:
    - 32 bytes = 256 bits of entropy
    - Base64 urlsafe encoding (uses - and _ instead of + and /)
    - No padding (= characters removed)
    - Example: "dGhpcyBpcyBhIDMyIGJ5dGUgbWFzdGVyIGtleQ"

    Security:
    - Uses os.urandom which is cryptographically secure on all platforms
    - 256 bits provides strong security against brute-force attacks
    - Each call produces a unique key

    Usage:
        key = generate_master_key()
        # Save to file or environment variable
        # Never commit to version control!
    """
    random_bytes = os.urandom(32)
    key_b64 = base64.urlsafe_b64encode(random_bytes).decode("ascii").rstrip("=")
    return key_b64
