"""API Key 加密存储与掩码展示。"""

from __future__ import annotations

import base64
import hashlib
import os

from cryptography.fernet import Fernet, InvalidToken

from . import config


def _load_key() -> bytes:
    """优先使用环境变量，其次复用本地密钥文件，都没有则生成新的。"""
    env = os.getenv("AI_IMAGE_SECRET_KEY")
    if env:
        return env.encode("utf-8")
    if config.SECRET_PATH.exists():
        return config.SECRET_PATH.read_bytes().strip()
    key = Fernet.generate_key()
    config.SECRET_PATH.parent.mkdir(parents=True, exist_ok=True)
    config.SECRET_PATH.write_bytes(key)
    return key


def _fernet() -> Fernet:
    raw = _load_key()
    try:
        return Fernet(raw)
    except Exception:
        derived = base64.urlsafe_b64encode(hashlib.sha256(raw).digest())
        return Fernet(derived)


def encrypt(plain: str | None) -> str:
    if not plain:
        return ""
    return _fernet().encrypt(plain.encode("utf-8")).decode("ascii")


def decrypt(token: str | None) -> str:
    if not token:
        return ""
    try:
        return _fernet().decrypt(token.encode("ascii")).decode("utf-8")
    except (InvalidToken, ValueError):
        return ""


def mask(plain: str | None) -> str:
    """展示用掩码：sk-a****abcd。"""
    if not plain:
        return ""
    if len(plain) <= 8:
        return plain[:2] + "****"
    return f"{plain[:3]}****-{plain[-4:]}"