"""Cryptographic operations implementation using cryptography library.

Provides AES-GCM encryption for audio data with envelope encryption.
"""

from __future__ import annotations

import secrets
import struct
from collections.abc import Iterator
from pathlib import Path
from typing import TYPE_CHECKING, BinaryIO, Final

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

from noteflow.infrastructure.logging import get_logger
from noteflow.infrastructure.security.protocols import EncryptedChunk

if TYPE_CHECKING:
    from noteflow.infrastructure.security.keystore import InMemoryKeyStore, KeyringKeyStore

logger = get_logger(__name__)

# Constants
KEY_SIZE: Final[int] = 32  # 256-bit key
NONCE_SIZE: Final[int] = 12  # 96-bit nonce for AES-GCM
TAG_SIZE: Final[int] = 16  # 128-bit authentication tag
MIN_CHUNK_LENGTH: Final[int] = NONCE_SIZE + TAG_SIZE  # Minimum valid encrypted chunk

# File format magic number and version
FILE_MAGIC: Final[bytes] = b"NFAE"  # NoteFlow Audio Encrypted
FILE_VERSION: Final[int] = 1


def _read_exact(handle: BinaryIO, size: int, description: str) -> bytes:
    """Read exactly size bytes or raise ValueError.

    Args:
        handle: File handle to read from.
        size: Number of bytes to read.
        description: Description for error message.

    Returns:
        Exactly size bytes.

    Raises:
        ValueError: If fewer than size bytes available.
    """
    data = handle.read(size)
    if len(data) < size:
        raise ValueError(f"Truncated {description}: expected {size}, got {len(data)}")
    return data


class AesGcmCryptoBox:
    """AES-GCM based encryption with envelope encryption.

    Uses a master key to wrap/unwrap per-meeting Data Encryption Keys (DEKs).
    Each audio chunk is encrypted with AES-256-GCM using the DEK.
    """

    def __init__(self, keystore: KeyringKeyStore | InMemoryKeyStore) -> None:
        """Initialize the crypto box.

        Args:
            keystore: KeyStore instance for master key access.
        """
        self._keystore = keystore
        self._master_cipher: AESGCM | None = None

    def _get_master_cipher(self) -> AESGCM:
        """Get or create the master key cipher."""
        if self._master_cipher is None:
            master_key = self._keystore.get_or_create_master_key()
            self._master_cipher = AESGCM(master_key)
        return self._master_cipher

    def generate_dek(self) -> bytes:
        """Generate a new Data Encryption Key.

        Returns:
            32-byte random DEK.
        """
        return secrets.token_bytes(KEY_SIZE)

    def wrap_dek(self, dek: bytes) -> bytes:
        """Encrypt DEK with master key.

        Args:
            dek: Data Encryption Key to wrap.

        Returns:
            Encrypted DEK (nonce || ciphertext || tag).
        """
        cipher = self._get_master_cipher()
        nonce = secrets.token_bytes(NONCE_SIZE)
        ciphertext = cipher.encrypt(nonce, dek, associated_data=None)
        # Return nonce || ciphertext (tag is appended by AESGCM)
        return nonce + ciphertext

    def unwrap_dek(self, wrapped_dek: bytes) -> bytes:
        """Decrypt DEK with master key.

        Args:
            wrapped_dek: Encrypted DEK from wrap_dek().

        Returns:
            Original DEK.

        Raises:
            ValueError: If decryption fails.
        """
        if len(wrapped_dek) < NONCE_SIZE + KEY_SIZE + TAG_SIZE:
            raise ValueError("Invalid wrapped DEK: too short")

        cipher = self._get_master_cipher()
        nonce = wrapped_dek[:NONCE_SIZE]
        ciphertext = wrapped_dek[NONCE_SIZE:]

        try:
            return cipher.decrypt(nonce, ciphertext, associated_data=None)
        except InvalidTag as e:
            raise ValueError(f"DEK unwrap failed: {e}") from e

    def encrypt_chunk(self, plaintext: bytes, dek: bytes) -> EncryptedChunk:
        """Encrypt a chunk of data with AES-GCM.

        Args:
            plaintext: Data to encrypt.
            dek: Data Encryption Key.

        Returns:
            EncryptedChunk with nonce, ciphertext, and tag.
        """
        cipher = AESGCM(dek)
        nonce = secrets.token_bytes(NONCE_SIZE)

        # AESGCM appends the tag to ciphertext
        ciphertext_with_tag = cipher.encrypt(nonce, plaintext, associated_data=None)

        # Split ciphertext and tag
        ciphertext = ciphertext_with_tag[:-TAG_SIZE]
        tag = ciphertext_with_tag[-TAG_SIZE:]

        return EncryptedChunk(nonce=nonce, ciphertext=ciphertext, tag=tag)

    def decrypt_chunk(self, chunk: EncryptedChunk, dek: bytes) -> bytes:
        """Decrypt a chunk of data.

        Args:
            chunk: EncryptedChunk to decrypt.
            dek: Data Encryption Key.

        Returns:
            Original plaintext.

        Raises:
            ValueError: If decryption fails.
        """
        cipher = AESGCM(dek)

        # Reconstruct ciphertext with tag for AESGCM
        ciphertext_with_tag = chunk.ciphertext + chunk.tag

        try:
            return cipher.decrypt(chunk.nonce, ciphertext_with_tag, associated_data=None)
        except InvalidTag as e:
            raise ValueError(f"Chunk decryption failed: {e}") from e


class ChunkedAssetWriter:
    """Streaming encrypted asset writer.

    File format:
    - 4 bytes: magic ("NFAE")
    - 1 byte: version
    - For each chunk:
        - 4 bytes: chunk length (big-endian)
        - 12 bytes: nonce
        - N bytes: ciphertext
        - 16 bytes: tag
    """

    def __init__(self, crypto: AesGcmCryptoBox) -> None:
        """Initialize the writer.

        Args:
            crypto: CryptoBox instance for encryption.
        """
        self._crypto = crypto
        self._file: Path | None = None
        self._dek: bytes | None = None
        self._handle: BinaryIO | None = None
        self._bytes_written: int = 0

    def open(self, path: Path, dek: bytes) -> None:
        """Open file for writing.

        Args:
            path: Path to the encrypted file.
            dek: Data Encryption Key for this file.
        """
        if self._handle is not None:
            raise RuntimeError("Already open")

        self._file = path
        self._dek = dek
        self._handle = path.open("wb")
        self._bytes_written = 0

        # Write header
        self._handle.write(FILE_MAGIC)
        self._handle.write(struct.pack("B", FILE_VERSION))

        logger.debug("Opened encrypted file for writing: %s", path)

    def write_chunk(self, audio_bytes: bytes) -> None:
        """Write and encrypt an audio chunk."""
        if self._handle is None or self._dek is None:
            raise RuntimeError("File not open")

        # Encrypt the chunk
        chunk = self._crypto.encrypt_chunk(audio_bytes, self._dek)

        # Calculate total chunk size (nonce + ciphertext + tag)
        chunk_data = chunk.nonce + chunk.ciphertext + chunk.tag
        chunk_length = len(chunk_data)

        # Write length prefix and chunk data
        self._handle.write(struct.pack(">I", chunk_length))
        self._handle.write(chunk_data)
        self._handle.flush()

        self._bytes_written += 4 + chunk_length

    def close(self) -> None:
        """Finalize and close the file."""
        if self._handle is not None:
            self._handle.close()
            self._handle = None
            logger.debug("Closed encrypted file, wrote %d bytes", self._bytes_written)

        self._dek = None

    @property
    def is_open(self) -> bool:
        """Check if file is open for writing."""
        return self._handle is not None

    @property
    def bytes_written(self) -> int:
        """Total encrypted bytes written."""
        return self._bytes_written


class ChunkedAssetReader:
    """Streaming encrypted asset reader."""

    def __init__(self, crypto: AesGcmCryptoBox) -> None:
        """Initialize the reader.

        Args:
            crypto: CryptoBox instance for decryption.
        """
        self._crypto = crypto
        self._file: Path | None = None
        self._dek: bytes | None = None
        self._handle: BinaryIO | None = None

    def open(self, path: Path, dek: bytes) -> None:
        """Open file for reading."""
        if self._handle is not None:
            raise RuntimeError("Already open")

        self._file = path
        self._dek = dek
        self._handle = path.open("rb")

        # Read and validate header
        magic = self._handle.read(4)
        if magic != FILE_MAGIC:
            self._handle.close()
            self._handle = None
            raise ValueError(f"Invalid file format: expected {FILE_MAGIC!r}, got {magic!r}")

        try:
            version_bytes = _read_exact(self._handle, 1, "version header")
        except ValueError as e:
            self._handle.close()
            self._handle = None
            raise ValueError(f"Invalid file format: {e}") from e

        version = struct.unpack("B", version_bytes)[0]
        if version != FILE_VERSION:
            self._handle.close()
            self._handle = None
            raise ValueError(f"Unsupported file version: {version}")

        logger.debug("Opened encrypted file for reading: %s", path)

    def read_chunks(self) -> Iterator[bytes]:
        """Yield decrypted audio chunks."""
        if self._handle is None or self._dek is None:
            raise RuntimeError("File not open")

        while True:
            # Read chunk length
            length_bytes = self._handle.read(4)
            if len(length_bytes) == 0:
                break  # Clean end of file
            if len(length_bytes) < 4:
                raise ValueError("Truncated chunk length header")

            chunk_length = struct.unpack(">I", length_bytes)[0]

            # Validate minimum chunk size (nonce + tag at minimum)
            if chunk_length < MIN_CHUNK_LENGTH:
                raise ValueError(
                    f"Invalid chunk length {chunk_length}: "
                    f"minimum is {MIN_CHUNK_LENGTH} (nonce + tag)"
                )

            # Read chunk data
            chunk_data = _read_exact(self._handle, chunk_length, "chunk data")

            # Parse chunk (nonce + ciphertext + tag)
            nonce = chunk_data[:NONCE_SIZE]
            ciphertext = chunk_data[NONCE_SIZE:-TAG_SIZE]
            tag = chunk_data[-TAG_SIZE:]

            chunk = EncryptedChunk(nonce=nonce, ciphertext=ciphertext, tag=tag)

            # Decrypt and yield
            yield self._crypto.decrypt_chunk(chunk, self._dek)

    def close(self) -> None:
        """Close the file."""
        if self._handle is not None:
            self._handle.close()
            self._handle = None
            logger.debug("Closed encrypted file")

        self._dek = None

    @property
    def is_open(self) -> bool:
        """Check if file is open for reading."""
        return self._handle is not None
