from __future__ import annotations

import base64
import json
import os
from collections.abc import Iterator
from pathlib import Path
from typing import Any

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

MAGIC = b"KIMUENC1"
NONCE_BYTES = 12
TAG_BYTES = 16
CHUNK_BYTES = 1024 * 1024
ENCRYPTION_VERSION = 1


class EncryptionError(RuntimeError):
    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code


def parse_master_key(value: str | bytes | None) -> bytes | None:
    if value is None or value == "":
        return None
    if isinstance(value, bytes):
        key = value
    else:
        raw = value.strip()
        try:
            key = base64.urlsafe_b64decode(raw + "=" * (-len(raw) % 4))
        except Exception as exc:  # noqa: BLE001 - normalise key errors
            raise EncryptionError(
                "KIMU_ENCRYPTION_KEY_INVALID",
                "KIMU_MASTER_KEY must be URL-safe base64 for exactly 32 bytes.",
            ) from exc
    if len(key) != 32:
        raise EncryptionError(
            "KIMU_ENCRYPTION_KEY_INVALID",
            "KIMU_MASTER_KEY must decode to exactly 32 bytes.",
        )
    return key


def generate_master_key() -> str:
    return base64.urlsafe_b64encode(os.urandom(32)).decode("ascii")


def file_aad(file_id: str, sha256: str, key_version: int) -> bytes:
    return f"kimu:file:v1:{key_version}:{file_id}:{sha256}".encode("utf-8")


def metadata_aad(file_id: str, sha256: str, key_version: int) -> bytes:
    return f"kimu:file-metadata:v1:{key_version}:{file_id}:{sha256}".encode("utf-8")


def encrypt_file(source: Path, destination: Path, key: bytes, aad: bytes) -> int:
    destination.parent.mkdir(parents=True, exist_ok=True)
    nonce = os.urandom(NONCE_BYTES)
    encryptor = Cipher(algorithms.AES(key), modes.GCM(nonce)).encryptor()
    encryptor.authenticate_additional_data(aad)
    temporary = destination.with_name(f".{destination.name}.encrypting")
    try:
        with source.open("rb") as reader, temporary.open("wb") as writer:
            writer.write(MAGIC)
            writer.write(nonce)
            for block in iter(lambda: reader.read(CHUNK_BYTES), b""):
                writer.write(encryptor.update(block))
            writer.write(encryptor.finalize())
            writer.write(encryptor.tag)
            writer.flush()
            os.fsync(writer.fileno())
        os.replace(temporary, destination)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise
    return destination.stat().st_size


def _header_and_lengths(path: Path) -> tuple[bytes, bytes, int]:
    total = path.stat().st_size
    minimum = len(MAGIC) + NONCE_BYTES + TAG_BYTES
    if total < minimum:
        raise EncryptionError("KIMU_ENCRYPTED_FILE_INVALID", "Encrypted file is truncated.")
    with path.open("rb") as handle:
        magic = handle.read(len(MAGIC))
        nonce = handle.read(NONCE_BYTES)
        handle.seek(-TAG_BYTES, os.SEEK_END)
        tag = handle.read(TAG_BYTES)
    if magic != MAGIC or len(nonce) != NONCE_BYTES or len(tag) != TAG_BYTES:
        raise EncryptionError("KIMU_ENCRYPTED_FILE_INVALID", "Encrypted file header is invalid.")
    ciphertext_bytes = total - len(MAGIC) - NONCE_BYTES - TAG_BYTES
    return nonce, tag, ciphertext_bytes


def iter_decrypted_file(path: Path, key: bytes, aad: bytes) -> Iterator[bytes]:
    nonce, tag, remaining = _header_and_lengths(path)
    decryptor = Cipher(algorithms.AES(key), modes.GCM(nonce, tag)).decryptor()
    decryptor.authenticate_additional_data(aad)
    try:
        with path.open("rb") as reader:
            reader.seek(len(MAGIC) + NONCE_BYTES)
            while remaining:
                block = reader.read(min(CHUNK_BYTES, remaining))
                if not block:
                    raise EncryptionError(
                        "KIMU_ENCRYPTED_FILE_INVALID", "Encrypted file ended unexpectedly."
                    )
                remaining -= len(block)
                plain = decryptor.update(block)
                if plain:
                    yield plain
            final = decryptor.finalize()
            if final:
                yield final
    except InvalidTag as exc:
        raise EncryptionError(
            "KIMU_ENCRYPTED_FILE_AUTH_FAILED",
            "Encrypted file authentication failed. The body or key may have changed.",
        ) from exc


def decrypt_file_bytes(
    path: Path, key: bytes, aad: bytes, maximum_bytes: int | None = None
) -> bytes:
    chunks: list[bytes] = []
    total = 0
    for chunk in iter_decrypted_file(path, key, aad):
        total += len(chunk)
        if maximum_bytes is not None and total > maximum_bytes:
            raise EncryptionError(
                "KIMU_DECRYPT_LIMIT",
                "The decrypted body exceeds the permitted in-memory limit.",
            )
        chunks.append(chunk)
    return b"".join(chunks)


def encrypt_json(value: dict[str, Any], key: bytes, aad: bytes) -> bytes:
    nonce = os.urandom(NONCE_BYTES)
    payload = json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    return nonce + AESGCM(key).encrypt(nonce, payload, aad)


def decrypt_json(value: bytes | memoryview | None, key: bytes, aad: bytes) -> dict[str, Any]:
    if value is None:
        return {}
    raw = bytes(value)
    if len(raw) <= NONCE_BYTES + TAG_BYTES:
        raise EncryptionError(
            "KIMU_ENCRYPTED_METADATA_INVALID", "Encrypted file metadata is truncated."
        )
    nonce, ciphertext = raw[:NONCE_BYTES], raw[NONCE_BYTES:]
    try:
        plain = AESGCM(key).decrypt(nonce, ciphertext, aad)
        decoded = json.loads(plain.decode("utf-8"))
    except (InvalidTag, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise EncryptionError(
            "KIMU_ENCRYPTED_METADATA_AUTH_FAILED",
            "Encrypted file metadata authentication failed.",
        ) from exc
    if not isinstance(decoded, dict):
        raise EncryptionError(
            "KIMU_ENCRYPTED_METADATA_INVALID", "Encrypted file metadata is invalid."
        )
    return decoded


def is_encrypted_file(path: Path) -> bool:
    try:
        with path.open("rb") as handle:
            return handle.read(len(MAGIC)) == MAGIC
    except OSError:
        return False