# 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.
This module is device-agnostic and can be used on any platform.
"""

import os
import struct
from typing import Generator

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):
        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
            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)

    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):
        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
                max_chunk_size = CHUNK_SIZE + TAG_SIZE + 1024
                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