# SPDX-License-Identifier: Apache-2.0
"""
Cryptographic operations for model encryption and decryption.

Uses AES-256-GCM for symmetric encryption with streaming support for large files.
"""

import os
import struct
from typing import BinaryIO, Generator, Optional, Tuple
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend

# File format constants
MAGIC_BYTES = b"VLLMENC\x00"
VERSION = 1
ALGORITHM_AES_256_GCM = 1
HEADER_SIZE = 32
CHUNK_SIZE = 64 * 1024 * 1024  # 64MB chunks for streaming
NONCE_SIZE = 12
TAG_SIZE = 16


def generate_dek() -> bytes:
    """Generate a random 256-bit Data Encryption Key (DEK)."""
    return AESGCM.generate_key(bit_length=256)


def compute_file_hash(file_path: str) -> bytes:
    """Compute SHA-256 hash of a file."""
    digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
    with open(file_path, 'rb') as f:
        while chunk := f.read(8192):
            digest.update(chunk)
    return digest.finalize()


class ModelEncryptor:
    """Encrypts model files using AES-256-GCM with streaming."""
    
    def __init__(self, dek: bytes):
        """
        Initialize encryptor with a Data Encryption Key.
        
        Args:
            dek: 256-bit AES key
        """
        if len(dek) != 32:
            raise ValueError("DEK must be 32 bytes (256 bits)")
        self.dek = dek
        self.aesgcm = AESGCM(dek)
    
    def encrypt_file(
        self,
        input_path: str,
        output_path: str,
        key_id: bytes,
    ) -> None:
        """
        Encrypt a model file.
        
        Args:
            input_path: Path to plaintext model file
            output_path: Path to write encrypted file
            key_id: 16-byte UUID identifying the encryption key
        """
        if len(key_id) != 16:
            raise ValueError("key_id must be 16 bytes (UUID)")
        
        with open(input_path, 'rb') as infile, \
             open(output_path, 'wb') as outfile:
            
            # Write header
            header = self._create_header(key_id)
            outfile.write(header)
            
            # Encrypt and write chunks
            chunk_num = 0
            while True:
                chunk = infile.read(CHUNK_SIZE)
                if not chunk:
                    break
                
                # Generate unique nonce for this chunk
                nonce = os.urandom(NONCE_SIZE)
                
                # Encrypt chunk (GCM mode produces ciphertext + tag)
                ciphertext = self.aesgcm.encrypt(nonce, chunk, None)
                
                # Write: chunk_size (4B) | nonce (12B) | ciphertext+tag
                chunk_header = struct.pack('<I', len(ciphertext))
                outfile.write(chunk_header + nonce + ciphertext)
                
                chunk_num += 1
    
    def _create_header(self, key_id: bytes) -> bytes:
        """Create encrypted file header."""
        header = bytearray(HEADER_SIZE)
        header[0:8] = MAGIC_BYTES
        header[8] = VERSION
        header[9] = ALGORITHM_AES_256_GCM
        header[10:26] = key_id
        # bytes 26-32 reserved for future use
        return bytes(header)


class ModelDecryptor:
    """Decrypts model files using AES-256-GCM with streaming."""
    
    def __init__(self, dek: bytes):
        """
        Initialize decryptor with a Data Encryption Key.
        
        Args:
            dek: 256-bit AES key
        """
        if len(dek) != 32:
            raise ValueError("DEK must be 32 bytes (256 bits)")
        self.dek = dek
        self.aesgcm = AESGCM(dek)
    
    def decrypt_file_stream(
        self,
        encrypted_path: str,
    ) -> Generator[bytes, None, None]:
        """
        Decrypt a model file in streaming fashion.
        
        Args:
            encrypted_path: Path to encrypted model file
            
        Yields:
            Decrypted chunks of data
            
        Raises:
            ValueError: If file format is invalid
            Exception: If decryption fails (wrong key or corrupted data)
        """
        with open(encrypted_path, 'rb') as f:
            # Read and validate header
            header = f.read(HEADER_SIZE)
            if len(header) != HEADER_SIZE:
                raise ValueError("Invalid encrypted file: header too short")
            
            self._validate_header(header)
            
            # Stream decrypt chunks
            while True:
                # Read chunk header
                chunk_header = f.read(4)
                if not chunk_header:
                    break
                
                if len(chunk_header) != 4:
                    raise ValueError("Invalid encrypted file: incomplete chunk header")
                
                chunk_size = struct.unpack('<I', chunk_header)[0]

                # Validate chunk_size before allocating to prevent OOM from malicious files
                max_chunk_size = CHUNK_SIZE + TAG_SIZE + 1024  # slight headroom
                if chunk_size > max_chunk_size:
                    raise ValueError(
                        f"Invalid encrypted file: chunk_size {chunk_size} exceeds "
                        f"maximum allowed {max_chunk_size}"
                    )

                # Read nonce
                nonce = f.read(NONCE_SIZE)
                if len(nonce) != NONCE_SIZE:
                    raise ValueError("Invalid encrypted file: incomplete nonce")
                
                # Read ciphertext + tag
                ciphertext = f.read(chunk_size)
                if len(ciphertext) != chunk_size:
                    raise ValueError("Invalid encrypted file: incomplete ciphertext")
                
                # Decrypt chunk
                try:
                    plaintext = self.aesgcm.decrypt(nonce, ciphertext, None)
                    yield plaintext
                except Exception as e:
                    raise Exception(f"Decryption failed: {e}. Wrong key or corrupted data.")
    
    def decrypt_file(
        self,
        encrypted_path: str,
        output_path: str,
    ) -> None:
        """
        Decrypt a model file to disk.
        
        Args:
            encrypted_path: Path to encrypted model file
            output_path: Path to write decrypted file
        """
        with open(output_path, 'wb') as outfile:
            for chunk in self.decrypt_file_stream(encrypted_path):
                outfile.write(chunk)
    
    def get_key_id(self, encrypted_path: str) -> bytes:
        """Extract key_id from encrypted file header."""
        with open(encrypted_path, 'rb') as f:
            header = f.read(HEADER_SIZE)
            if len(header) != HEADER_SIZE:
                raise ValueError("Invalid encrypted file: header too short")
            self._validate_header(header)
            return header[10:26]
    
    def _validate_header(self, header: bytes) -> None:
        """Validate encrypted file header."""
        if header[0:8] != MAGIC_BYTES:
            raise ValueError("Invalid encrypted file: wrong magic bytes")
        
        version = header[8]
        if version != VERSION:
            raise ValueError(f"Unsupported file version: {version}")
        
        algorithm = header[9]
        if algorithm != ALGORITHM_AES_256_GCM:
            raise ValueError(f"Unsupported encryption algorithm: {algorithm}")


def is_encrypted_file(file_path: str) -> bool:
    """Check if a file is encrypted by reading its magic bytes."""
    try:
        with open(file_path, 'rb') as f:
            magic = f.read(8)
            return magic == MAGIC_BYTES
    except Exception:
        return False
