# -*- coding: utf-8 -*-
"""模型密钥的落库加密。

模型清单里存的是各家厂商的 API Key。这些密钥现在明文躺在每台客户端的
ai_settings.json 里——收进后端本来就是为了解决这个问题，如果后端再明文存一遍，
等于把泄漏面从"很多台客户端"换成"一台服务器 + 一个数据库备份文件"，并没有真的
变安全。

主密钥单独放一个 0600 的文件，不进数据库、不进代码库、不进备份脚本的默认范围。
数据库整个被拖走也解不开密钥。

用 AES-GCM：带认证标签，密文被改过会解密失败而不是悄悄返回一串垃圾——那种
"垃圾密钥"会变成一整天查不出来的 401。
"""

from __future__ import annotations

import base64
import os
import secrets
from pathlib import Path

_KEY_ENV = "WECOM_BACKEND_SECRET_KEY"
_KEY_FILENAME = "backend_secret.key"


class SecretBoxError(RuntimeError):
    pass


def _key_path(data_dir: Path | str) -> Path:
    return Path(data_dir) / _KEY_FILENAME


def load_or_create_key(data_dir: Path | str) -> bytes:
    """取主密钥。环境变量优先，方便容器部署不落盘。"""
    from_env = os.environ.get(_KEY_ENV, "").strip()
    if from_env:
        try:
            key = base64.urlsafe_b64decode(from_env + "=" * (-len(from_env) % 4))
        except Exception as exc:
            raise SecretBoxError(f"{_KEY_ENV} 不是合法的 base64") from exc
        if len(key) != 32:
            raise SecretBoxError(f"{_KEY_ENV} 解出来不是 32 字节")
        return key

    path = _key_path(data_dir)
    if path.exists():
        key = path.read_bytes().strip()
        key = base64.urlsafe_b64decode(key + b"=" * (-len(key) % 4))
        if len(key) != 32:
            raise SecretBoxError(f"{path} 里的主密钥长度不对，拒绝继续")
        return key

    key = secrets.token_bytes(32)
    path.parent.mkdir(parents=True, exist_ok=True)
    # 先按 0600 建文件再写内容，避免密钥在世界可读的瞬间落盘
    handle = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    try:
        os.write(handle, base64.urlsafe_b64encode(key))
    finally:
        os.close(handle)
    try:
        os.chmod(str(path), 0o600)
    except OSError:
        pass
    return key


def encrypt(plaintext: str, key: bytes) -> str:
    """→ base64(nonce | tag | 密文)。空串原样返回空串。"""
    text = str(plaintext or "")
    if not text:
        return ""
    from Crypto.Cipher import AES

    nonce = secrets.token_bytes(12)
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    body, tag = cipher.encrypt_and_digest(text.encode("utf-8"))
    return base64.urlsafe_b64encode(nonce + tag + body).decode("ascii")


def decrypt(token: str, key: bytes) -> str:
    """解不开就抛异常，绝不返回半截垃圾——那会变成查不出来的 401。"""
    raw = str(token or "")
    if not raw:
        return ""
    from Crypto.Cipher import AES

    try:
        blob = base64.urlsafe_b64decode(raw + "=" * (-len(raw) % 4))
        nonce, tag, body = blob[:12], blob[12:28], blob[28:]
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
        return cipher.decrypt_and_verify(body, tag).decode("utf-8")
    except Exception as exc:
        raise SecretBoxError("模型密钥解密失败：主密钥不匹配或密文被改动") from exc


def masked(plaintext: str) -> str:
    """给界面看的遮罩。永远不回显完整密钥。"""
    text = str(plaintext or "")
    if not text:
        return ""
    if len(text) <= 8:
        return "*" * len(text)
    return f"{text[:4]}{'*' * 6}{text[-4:]}"