"""AES-256-GCM 加解密，与 aiclient Tauri 端一致。"""

from __future__ import annotations

import base64
import json
import os
from typing import Any

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

ENVELOPE_VERSION = 1
ENCRYPTED_HEADER = "x-aiclient-encrypted"


def parse_aes_key(raw: str) -> bytes | None:
    """解析 32 字节密钥：标准 Base64 或 64 位十六进制。"""
    s = raw.strip()
    if not s:
        return None
    try:
        if len(s) == 64 and all(c in "0123456789abcdefABCDEF" for c in s):
            key = bytes.fromhex(s)
        else:
            key = base64.b64decode(s)
        if len(key) != 32:
            return None
        return key
    except Exception:
        return None


def encrypt_bytes(plaintext: bytes, key: bytes) -> str:
    nonce = os.urandom(12)
    ciphertext = AESGCM(key).encrypt(nonce, plaintext, None)
    return base64.b64encode(nonce + ciphertext).decode("ascii")


def decrypt_payload(b64_payload: str, key: bytes) -> bytes:
    raw = base64.b64decode(b64_payload)
    if len(raw) < 13:
        raise ValueError("密文过短")
    nonce, ciphertext = raw[:12], raw[12:]
    return AESGCM(key).decrypt(nonce, ciphertext, None)


def wrap_encrypted(plaintext: bytes, key: bytes) -> bytes:
    payload = encrypt_bytes(plaintext, key)
    return json.dumps({"v": ENVELOPE_VERSION, "payload": payload}, ensure_ascii=False).encode(
        "utf-8"
    )


def unwrap_encrypted(body: bytes, key: bytes) -> bytes:
    obj: dict[str, Any] = json.loads(body)
    if obj.get("v") != ENVELOPE_VERSION or "payload" not in obj:
        raise ValueError("无效的加密信封")
    payload = obj["payload"]
    if not isinstance(payload, str):
        raise ValueError("payload 应为字符串")
    return decrypt_payload(payload, key)


def is_encrypted_envelope(body: bytes) -> bool:
    try:
        obj = json.loads(body)
        return isinstance(obj, dict) and "payload" in obj and obj.get("v") == ENVELOPE_VERSION
    except Exception:
        return False