from __future__ import annotations

import base64
import hashlib
import os
import sys
from pathlib import Path
from typing import Iterable

try:
    from cryptography.exceptions import InvalidTag
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
except Exception as exc:  # noqa: BLE001
    AESGCM = None  # type: ignore[assignment]
    InvalidTag = Exception  # type: ignore[assignment]
    _CRYPTO_IMPORT_ERROR = exc
else:
    _CRYPTO_IMPORT_ERROR = None


MAGIC = b"TQFWENC1"
NONCE_SIZE = 12
SHA256_SIZE = 32
MIN_PAYLOAD_SIZE = len(MAGIC) + NONCE_SIZE + SHA256_SIZE + 16
DEFAULT_FIRMWARE_KEY_RELATIVE_PATH = Path("keys") / "firmware" / "fw_key.bin"


class FirmwareEnvelopeError(Exception):
    pass


def _require_crypto() -> None:
    if _CRYPTO_IMPORT_ERROR is not None or AESGCM is None:
        raise FirmwareEnvelopeError("cryptography 不可用，无法处理 .enc 固件。")


def _parse_key_material(raw: bytes, source: str) -> bytes:
    if len(raw) == 32:
        return raw

    try:
        text = raw.decode("utf-8").strip()
    except UnicodeDecodeError as exc:
        raise FirmwareEnvelopeError(f"固件密钥格式无效（{source}）：既不是 32 字节原始密钥，也不是文本编码。") from exc

    if not text:
        raise FirmwareEnvelopeError(f"固件密钥为空（{source}）。")

    lowered = text.lower()
    if lowered.startswith("hex:"):
        text = text[4:].strip()
    if len(text) == 64:
        try:
            key = bytes.fromhex(text)
        except ValueError as exc:
            raise FirmwareEnvelopeError(f"固件密钥 HEX 无效（{source}）。") from exc
        if len(key) == 32:
            return key

    if lowered.startswith("base64:"):
        text = text[7:].strip()
    try:
        key = base64.b64decode(text.encode("utf-8"), validate=True)
    except Exception:  # noqa: BLE001
        key = b""
    if len(key) == 32:
        return key

    raise FirmwareEnvelopeError(
        f"固件密钥长度无效（{source}）：需要 32 字节，或 64 位 HEX，或 32 字节 Base64。"
    )


def _iter_runtime_roots() -> Iterable[Path]:
    exe_path = Path(sys.executable).resolve()
    yield exe_path.parent
    for parent in exe_path.parents:
        if parent.suffix.lower() == ".app":
            yield parent.parent
            break


def _candidate_key_paths() -> Iterable[Path]:
    if getattr(sys, "frozen", False):
        for root in _iter_runtime_roots():
            yield root / DEFAULT_FIRMWARE_KEY_RELATIVE_PATH

    for parent in Path(__file__).resolve().parents:
        yield parent / DEFAULT_FIRMWARE_KEY_RELATIVE_PATH


def _load_key_from_fixed_path(purpose: str) -> bytes:
    for path in _candidate_key_paths():
        if not path.exists():
            continue
        try:
            raw = path.read_bytes()
        except OSError as exc:
            raise FirmwareEnvelopeError(f"{purpose}密钥文件读取失败：{path}") from exc
        return _parse_key_material(raw, str(path))

    raise FirmwareEnvelopeError(
        f"缺少{purpose}密钥文件。请在运行目录提供 {DEFAULT_FIRMWARE_KEY_RELATIVE_PATH}。"
    )


def load_firmware_encrypt_key() -> bytes:
    return _load_key_from_fixed_path(purpose="固件加密")


def load_firmware_decrypt_key() -> bytes:
    return _load_key_from_fixed_path(purpose="固件解密")


def encrypt_firmware_blob(plain: bytes, key: bytes) -> bytes:
    _require_crypto()
    if len(key) != 32:
        raise FirmwareEnvelopeError("固件加密密钥长度无效，必须为 32 字节。")
    nonce = os.urandom(NONCE_SIZE)
    digest = hashlib.sha256(plain).digest()
    aes = AESGCM(key)
    ciphertext = aes.encrypt(nonce, plain, MAGIC)
    return MAGIC + nonce + digest + ciphertext


def decrypt_firmware_blob(payload: bytes, key: bytes) -> bytes:
    _require_crypto()
    if len(key) != 32:
        raise FirmwareEnvelopeError("固件解密密钥长度无效，必须为 32 字节。")
    if len(payload) < MIN_PAYLOAD_SIZE:
        raise FirmwareEnvelopeError("固件 .enc 文件格式无效：长度不足。")
    if payload[: len(MAGIC)] != MAGIC:
        raise FirmwareEnvelopeError("固件 .enc 文件格式无效：magic 不匹配。")

    cursor = len(MAGIC)
    nonce = payload[cursor : cursor + NONCE_SIZE]
    cursor += NONCE_SIZE
    expected_digest = payload[cursor : cursor + SHA256_SIZE]
    cursor += SHA256_SIZE
    ciphertext = payload[cursor:]

    aes = AESGCM(key)
    try:
        plain = aes.decrypt(nonce, ciphertext, MAGIC)
    except InvalidTag as exc:
        raise FirmwareEnvelopeError("固件 .enc 解密失败：密钥错误或密文损坏。") from exc
    except Exception as exc:  # noqa: BLE001
        raise FirmwareEnvelopeError("固件 .enc 解密失败。") from exc

    digest = hashlib.sha256(plain).digest()
    if digest != expected_digest:
        raise FirmwareEnvelopeError("固件 .enc 解密后哈希不一致。")
    return plain