#!/usr/bin/env python3
"""AES-256-GCM encryption, key derivation, and LUKS container management."""

import os
import struct
import subprocess
import hashlib
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

try:
    from argon2.low_level import hash_secret_raw, Type
    HAS_ARGON2 = True
except ImportError:
    HAS_ARGON2 = False

NONCE_SIZE = 12
KEY_SIZE = 32
SALT_SIZE = 16
HEADER_MAGIC = b"BB01"
HEADER_FMT = f"!4s{SALT_SIZE}s{NONCE_SIZE}s"
HEADER_SIZE = struct.calcsize(HEADER_FMT)


def derive_key(password: bytes, salt: bytes, method: str = "argon2id") -> bytes:
    if method == "argon2id" and HAS_ARGON2:
        return hash_secret_raw(
            secret=password,
            salt=salt,
            time_cost=3,
            memory_cost=65536,
            parallelism=4,
            hash_len=KEY_SIZE,
            type=Type.ID,
        )
    return PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=salt,
        iterations=600_000,
    ).derive(password)


def derive_network_key(network_secret: bytes, device_serial: bytes) -> bytes:
    return HKDF(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=device_serial,
        info=b"bigbrother-luks-unlock",
    ).derive(network_secret)


class CryptoEngine:
    def __init__(self, key: bytes):
        if len(key) != KEY_SIZE:
            raise ValueError(f"Key must be {KEY_SIZE} bytes")
        self._aesgcm = AESGCM(key)
        self._key = key

    def encrypt(self, plaintext: bytes, aad: Optional[bytes] = None) -> bytes:
        nonce = os.urandom(NONCE_SIZE)
        ct = self._aesgcm.encrypt(nonce, plaintext, aad)
        return nonce + ct

    def decrypt(self, data: bytes, aad: Optional[bytes] = None) -> bytes:
        nonce = data[:NONCE_SIZE]
        ct = data[NONCE_SIZE:]
        return self._aesgcm.decrypt(nonce, ct, aad)


def encrypt_file(
    src: str,
    dst: str,
    password: bytes,
    method: str = "argon2id",
    chunk_size: int = 64 * 1024,
) -> None:
    salt = os.urandom(SALT_SIZE)
    key = derive_key(password, salt, method)
    engine = CryptoEngine(key)

    with open(src, "rb") as fin:
        plaintext = fin.read()

    nonce = os.urandom(NONCE_SIZE)
    aesgcm = AESGCM(key)
    ct = aesgcm.encrypt(nonce, plaintext, None)

    with open(dst, "wb") as fout:
        fout.write(struct.pack(HEADER_FMT, HEADER_MAGIC, salt, nonce))
        fout.write(ct)


def decrypt_file(
    src: str,
    dst: str,
    password: bytes,
    method: str = "argon2id",
) -> None:
    with open(src, "rb") as fin:
        header = fin.read(HEADER_SIZE)
        magic, salt, nonce = struct.unpack(HEADER_FMT, header)
        if magic != HEADER_MAGIC:
            raise ValueError("Invalid encrypted file header")
        ct = fin.read()

    key = derive_key(password, salt, method)
    aesgcm = AESGCM(key)
    plaintext = aesgcm.decrypt(nonce, ct, None)

    with open(dst, "wb") as fout:
        fout.write(plaintext)


class LUKSContainer:
    def __init__(self, container_path: str, mount_point: str, mapper_name: str = "bb_storage"):
        self.container_path = container_path
        self.mount_point = mount_point
        self.mapper_name = mapper_name
        self.device_path = f"/dev/mapper/{mapper_name}"

    def create(self, size_mb: int, passphrase: bytes) -> bool:
        try:
            subprocess.run(
                ["dd", "if=/dev/zero", f"of={self.container_path}",
                 "bs=1M", f"count={size_mb}"],
                check=True, capture_output=True,
            )
            proc = subprocess.run(
                ["cryptsetup", "luksFormat", "--type", "luks2",
                 "--cipher", "aes-xts-plain64", "--key-size", "512",
                 "--hash", "sha256", "--batch-mode", self.container_path],
                input=passphrase + b"\n",
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    def open(self, passphrase: bytes) -> bool:
        try:
            subprocess.run(
                ["cryptsetup", "open", self.container_path,
                 self.mapper_name, "--type", "luks2"],
                input=passphrase + b"\n",
                check=True, capture_output=True,
            )
            Path(self.mount_point).mkdir(parents=True, exist_ok=True)
            subprocess.run(
                ["mount", self.device_path, self.mount_point],
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    def close(self) -> bool:
        try:
            subprocess.run(
                ["umount", self.mount_point],
                check=True, capture_output=True,
            )
            subprocess.run(
                ["cryptsetup", "close", self.mapper_name],
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    def destroy_header(self) -> bool:
        try:
            subprocess.run(
                ["dd", "if=/dev/urandom", f"of={self.container_path}",
                 "bs=1M", "count=2", "conv=notrunc"],
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    @property
    def is_open(self) -> bool:
        return Path(self.device_path).exists()