"""Manifest encryption at rest — AES-256-GCM with a compact binary format.

File format::

    CENSUS_ENC\\x01     (10 bytes: magic + version)
    <nonce: 12 bytes>
    <ciphertext + GCM tag>

The ``cryptography`` library is a transitive dependency (via the
CertiSigma SDK), so no additional install is required.

Security properties:
- 256-bit AES in GCM mode → authenticated encryption (confidentiality + integrity)
- 96-bit random nonce per file (NIST SP 800-38D § 8.2)
- GCM tag is appended by AESGCM (16 bytes)
- Keys are 64 hex characters (32 bytes / 256 bits)
"""

from __future__ import annotations

import logging
import os
import re
from pathlib import Path

logger = logging.getLogger(__name__)

MAGIC = b"CENSUS_ENC"
FORMAT_VERSION = b"\x01"
HEADER = MAGIC + FORMAT_VERSION  # 11 bytes
NONCE_SIZE = 12  # 96-bit nonce per NIST SP 800-38D

_KEY_RE = re.compile(r"^[0-9a-fA-F]{64}$")


def validate_key(key_hex: str) -> bytes:
    """Validate a hex-encoded AES-256 key and return raw bytes.

    Raises ``ValueError`` on invalid format.
    """
    if not isinstance(key_hex, str) or not _KEY_RE.match(key_hex):
        raise ValueError(
            "Encryption key must be exactly 64 hex characters (256 bits). "
            "Generate one with: census key-gen"
        )
    return bytes.fromhex(key_hex)


def encrypt_file(src: Path, dst: Path, key_hex: str) -> None:
    """Encrypt *src* to *dst* using AES-256-GCM.

    The output file uses the Census encrypted format (magic header + nonce +
    ciphertext). The source file is read entirely into memory — acceptable
    for SQLite manifests (typically < 100 MB).

    *dst* is written atomically: a temp file is created alongside, then
    renamed on success.
    """
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    key = validate_key(key_hex)
    plaintext = src.read_bytes()
    nonce = os.urandom(NONCE_SIZE)
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext, associated_data=HEADER)

    tmp_path = dst.with_suffix(dst.suffix + ".tmp")
    fd = os.open(str(tmp_path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    try:
        with os.fdopen(fd, "wb") as f:
            f.write(HEADER)
            f.write(nonce)
            f.write(ciphertext)
        os.rename(str(tmp_path), str(dst))
    except BaseException:
        try:
            os.unlink(str(tmp_path))
        except OSError as exc:
            logger.warning("Cannot remove temporary encryption file %s: %s", tmp_path, exc)
        raise


def decrypt_file(src: Path, key_hex: str) -> bytes:
    """Decrypt a Census encrypted file and return the plaintext bytes.

    Raises ``ValueError`` on format errors, ``cryptography.exceptions.InvalidTag``
    on wrong key or tampered data.
    """
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    key = validate_key(key_hex)
    data = src.read_bytes()

    header_len = len(HEADER)
    if len(data) < header_len + NONCE_SIZE:
        raise ValueError(f"File too small to be a Census encrypted file: {src}")
    if data[:header_len] != HEADER:
        raise ValueError(f"Not a Census encrypted file (bad magic): {src}")

    nonce = data[header_len:header_len + NONCE_SIZE]
    ciphertext = data[header_len + NONCE_SIZE:]

    aesgcm = AESGCM(key)
    return aesgcm.decrypt(nonce, ciphertext, associated_data=HEADER)


def is_encrypted_file(path: Path) -> bool:
    """Return True if *path* starts with the Census encryption magic header."""
    if not path.exists():
        return False
    try:
        with path.open("rb") as f:
            header = f.read(len(HEADER))
            return header == HEADER
    except OSError:
        return False
