"""加密 DB 列 secrets store — AES-256-GCM (KTD8)。

KTD8 关键决策：
- 每行使用随机 96-bit nonce + HKDF with per-row salt 派生 per-row 密钥。
- 生产环境 master key 必须来自云 KMS；环境变量 AGENTKIT_MASTER_KEY 仅作开发 fallback。
- 应用启动 guard：生产模式下若仅检测到环境变量 key 应拒绝启动
  （由 server 启动钩子实现，本模块提供 ``assert_production_master_key`` 辅助）。
- Master key 轮换采用双密钥窗口策略（key_id 字段标记）。

U3：Redis 后端支持多 worker 共享。``redis`` 参数为 None 时降级到内存字典
（仅开发模式），生产环境必须注入 Redis 客户端。
"""

from __future__ import annotations

import base64
import logging
import os

import redis.asyncio as aioredis
from pydantic import BaseModel, ConfigDict

logger = logging.getLogger(__name__)

# AES-256-GCM 参数
NONCE_SIZE = 12  # 96-bit nonce (GCM 推荐长度)
KEY_SIZE = 32  # 256-bit key
SALT_SIZE = 16  # per-row salt
HKDF_INFO = b"agentkit-channels-secrets-v1"

# 生产环境标记：当此环境变量存在时启用生产 guard。
_PRODUCTION_ENV_VAR = "AGENTKIT_ENV"
_PRODUCTION_VALUE = "production"


class SecretEntry(BaseModel):
    """凭证条目 — 对应 DB 一行加密列。"""

    model_config = ConfigDict()

    key: str  # 唯一键（如 "feishu:app_id:xxx"）
    value: str  # base64 编码的密文
    nonce: str  # base64 编码的 96-bit nonce
    salt: str  # base64 编码的 per-row salt
    key_id: str = "default"  # master key ID（用于双密钥轮换窗口）
    created_at: str = ""
    updated_at: str = ""


def assert_production_master_key(master_key: bytes | None, *, source: str = "env") -> None:
    """生产环境启动 guard。

    若运行于生产模式（``AGENTKIT_ENV=production``）且 master key 来源于
    环境变量（source="env"）或为空，则拒绝启动。

    Args:
        master_key: 已加载的 master key 字节，可能为 None。
        source: master key 来源标记。"env" 表示环境变量 fallback，
            "kms" 表示云 KMS。生产模式仅接受 "kms"。

    Raises:
        RuntimeError: 生产模式下 master key 来源不可信。
    """
    if os.environ.get(_PRODUCTION_ENV_VAR) == _PRODUCTION_VALUE:
        if source != "kms" or not master_key:
            raise RuntimeError(
                "生产环境 master key 必须来自云 KMS（source=kms）；"
                "环境变量 AGENTKIT_MASTER_KEY 仅作开发 fallback。"
                "请通过 KMS 注入 master key 并以 source='kms' 调用本 guard。"
            )


class SecretsStore:
    """加密凭证存储。

    使用 AES-256-GCM 加密，HKDF with per-row salt 派生 per-row 密钥。
    Master key 从环境变量 ``AGENTKIT_MASTER_KEY`` 读取（开发 fallback）；
    生产环境应通过 KMS 提供 master key 并显式传入。

    U3：``redis`` 参数注入 Redis 客户端后，凭证存储在 Redis（多 worker 共享）。
    ``redis=None`` 时降级到内存字典（仅开发模式）。
    """

    # Redis key 前缀 — 遵循 codebase 约定（agentkit:<subsystem>:）
    _REDIS_PREFIX = "agentkit:secrets:"

    def __init__(
        self,
        master_key: bytes | None = None,
        *,
        key_source: str = "env",
        redis: aioredis.Redis | None = None,
    ):
        """初始化 secrets store。

        Args:
            master_key: 显式传入的 master key。若为 None 则从环境变量加载。
            key_source: master key 来源标记，传给 ``assert_production_master_key``。
                生产环境应设为 "kms"。
            redis: Redis 客户端（``redis.asyncio.Redis``）。注入后凭证存储在 Redis，
                多 worker 共享。None 时降级到内存字典（仅开发模式）。
        """
        self._master_key = master_key or self._load_master_key()
        # 生产 guard：若 key_source="env" 且在生产模式下，构造即失败。
        assert_production_master_key(self._master_key, source=key_source)
        self._redis = redis
        # 内存降级存储（redis=None 时使用）
        self._store: dict[str, SecretEntry] = {}
        if redis is None:
            logger.warning(
                "SecretsStore 运行于内存降级模式（redis=None）— "
                "多 worker 部署下状态不共享，仅限开发使用。"
            )

    def _load_master_key(self) -> bytes:
        """从环境变量加载 master key（开发 fallback）。"""
        key_b64 = os.environ.get("AGENTKIT_MASTER_KEY")
        if key_b64:
            return base64.b64decode(key_b64)
        # ponytail: 开发模式生成临时 key。生产环境必须通过 KMS 提供
        # （由 assert_production_master_key 在 __init__ 中强制）。
        logger.warning("AGENTKIT_MASTER_KEY 未设置 — 使用临时 key（仅限开发）")
        return os.urandom(KEY_SIZE)

    def _derive_key(self, salt: bytes) -> bytes:
        """HKDF 从 master key + per-row salt 派生 per-row 密钥。"""
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.kdf.hkdf import HKDF

        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=KEY_SIZE,
            salt=salt,
            info=HKDF_INFO,
        )
        return hkdf.derive(self._master_key)

    def encrypt(self, plaintext: str) -> SecretEntry:
        """加密凭证，返回 SecretEntry（key 字段留空，由调用方设置）。"""
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM

        salt = os.urandom(SALT_SIZE)
        key = self._derive_key(salt)
        nonce = os.urandom(NONCE_SIZE)

        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, plaintext.encode(), None)

        return SecretEntry(
            key="",
            value=base64.b64encode(ciphertext).decode(),
            nonce=base64.b64encode(nonce).decode(),
            salt=base64.b64encode(salt).decode(),
        )

    def decrypt(self, entry: SecretEntry) -> str:
        """解密凭证。密钥不匹配或密文损坏时抛出 cryptography 异常。"""
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM

        salt = base64.b64decode(entry.salt)
        key = self._derive_key(salt)
        nonce = base64.b64decode(entry.nonce)
        ciphertext = base64.b64decode(entry.value)

        aesgcm = AESGCM(key)
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
        return plaintext.decode()

    async def set_secret(self, key: str, value: str) -> SecretEntry:
        """存储加密凭证。覆盖同名 key。"""
        entry = self.encrypt(value)
        entry.key = key
        if self._redis is not None:
            await self._redis.set(self._REDIS_PREFIX + key, entry.model_dump_json())
        else:
            self._store[key] = entry
        return entry

    async def get_secret(self, key: str) -> str | None:
        """读取并解密凭证。key 不存在返回 None。"""
        if self._redis is not None:
            raw = await self._redis.get(self._REDIS_PREFIX + key)
            if raw is None:
                return None
            entry = SecretEntry.model_validate_json(raw)
        else:
            entry = self._store.get(key)
            if entry is None:
                return None
        return self.decrypt(entry)

    async def delete_secret(self, key: str) -> bool:
        """删除凭证。返回是否删除成功。"""
        if self._redis is not None:
            deleted = await self._redis.delete(self._REDIS_PREFIX + key)
            return deleted > 0
        if key in self._store:
            del self._store[key]
            return True
        return False

    async def list_keys(self, prefix: str | None = None) -> list[str]:
        """列出凭证键。可选前缀过滤。"""
        if self._redis is not None:
            pattern = self._REDIS_PREFIX + (prefix or "") + "*"
            keys: list[str] = []
            async for k in self._redis.scan_iter(match=pattern):
                keys.append(k[len(self._REDIS_PREFIX) :])
            return keys
        if prefix:
            return [k for k in self._store if k.startswith(prefix)]
        return list(self._store.keys())
