# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT

"""
AES 加密算法.
"""

__all__ = [
    "AES_KEY_LEN",
    "AES_MAC_LEN",
    "AES_NONCE_LEN",
    "AesKey",
    "FileMode",
]

import base64
import secrets
from dataclasses import dataclass
from io import BytesIO

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from typing import IO, Literal, Optional, Tuple
from typing_extensions import Self

from .. import error

AES_KEY_LEN = 32

AES_NONCE_LEN = 12

AES_MAC_LEN = 16

CHUNK_SIZE = 1024

FileMode = Literal["b", "t"]
"""
文件加密模式: b 表示二进制, t 表示文本 (按行加密, 加密数据 Base64 编码后按行保存).
"""


@dataclass(eq=False, frozen=True)
class AesKey:
    """
    AES 密钥.

    本类不含可变状态, 所有方法是线程安全的.
    """

    key: bytes

    def __post_init__(self):
        if len(self.key) != AES_KEY_LEN:
            raise error.ParamError("key", "Invalid AES key length")

    @classmethod
    def generate(cls) -> Self:
        """
        生成 AES 密钥.
        """
        key = secrets.token_bytes(AES_KEY_LEN)

        return cls(key)

    def encrypt(
        self,
        plaintext: bytes,
        nonce: Optional[bytes] = None,
    ) -> Tuple[bytes, bytes, bytes]:
        """
        用 AES-GCM 加密数据.

        Returns:
            单次随机数, 密文, 及校验码.

        Raises:
            EncryptionError:
        """
        with BytesIO(plaintext) as source, BytesIO() as dest:
            nonce, mac = self.encrypt_stream(source, dest, "b", nonce)
            ciphertext = dest.getvalue()

        return nonce, ciphertext, mac

    def encrypt_stream(
        self,
        source: IO,
        dest: IO,
        mode: FileMode,
        nonce: Optional[bytes] = None,
    ) -> Tuple[bytes, bytes]:
        """
        用 AES-GCM 加密数据流.

        Returns:
            单次随机数及校验码.

        Raises:
            EncryptionError:
        """
        if nonce is None:
            nonce = secrets.token_bytes(AES_NONCE_LEN)
        elif len(nonce) != AES_NONCE_LEN:
            raise error.ParamError("nonce", "wrong length")

        try:
            aes = Cipher(algorithms.AES256(self.key), modes.GCM(nonce)).encryptor()

            if mode == "b":
                source_b: IO[bytes] = source
                dest_b: IO[bytes] = dest
                while plaintext := source_b.read(CHUNK_SIZE):
                    ciphertext = aes.update(plaintext)
                    dest_b.write(ciphertext)
            else:
                source_t: IO[str] = source
                dest_t: IO[str] = dest
                for plaintext in source_t:
                    plaintext = plaintext.rstrip("\n")
                    ciphertext = aes.update(plaintext.encode())
                    dest_t.write(base64.b64encode(ciphertext).decode())
                    dest_t.write("\n")

            ciphertext = aes.finalize()
            mac = aes.tag
        except Exception as e:
            raise error.EncryptionError("AES encryption failed") from e

        return nonce, mac

    def decrypt(self, nonce: bytes, ciphertext: bytes, mac: Optional[bytes]) -> bytes:
        """
        用 AES-GCM 解密并校验数据.

        Raises:
            DecryptionError:
        """
        with BytesIO(ciphertext) as source, BytesIO() as dest:
            self.decrypt_stream(source, dest, nonce, mac, "b")
            plaintext = dest.getvalue()

        return plaintext

    def decrypt_stream(
        self,
        source: IO,
        dest: IO,
        nonce: bytes,
        mac: Optional[bytes],
        mode: FileMode,
    ) -> None:
        """
        用 AES-GCM 解密并校验数据流.

        Raises:
            DecryptionError:
        """
        try:
            aes = Cipher(algorithms.AES256(self.key), modes.GCM(nonce)).decryptor()

            if mode == "b":
                source_b: IO[bytes] = source
                dest_b: IO[bytes] = dest
                while ciphertext := source_b.read(CHUNK_SIZE):
                    plaintext = aes.update(ciphertext)
                    dest_b.write(plaintext)
            else:
                source_t: IO[str] = source
                dest_t: IO[str] = dest
                for ciphertext in source_t:
                    ciphertext = ciphertext.rstrip("\n")
                    plaintext = aes.update(base64.b64decode(ciphertext)).decode()
                    dest_t.write(plaintext)
                    dest_t.write("\n")

            if mac is not None:
                aes.finalize_with_tag(mac)
        except Exception as e:
            raise error.DecryptionError("AES decryption failed") from e
