#!/usr/bin/env python3
"""kb_encrypt_backfill.py —— 对 KB 副本的 registry JSONL 做 payload 级 AES-256-GCM 加密。

加密语义(方案 §五):
- 断言文本字段(assertion_text / evidence_text 等)与断言导出文件 → AES-256-GCM
- 向量(向量库存储)保持明文(检索数学需要)
- 图存储文件保持明文(存储引擎限制,包内)
- 本脚本只处理交付副本,绝不触碰生产库
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import sys
from pathlib import Path

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

NONCE_LEN = 12


def load_key(key_file: Path) -> bytes:
    key_b64 = key_file.read_text(encoding="utf-8").strip()
    key = base64.b64decode(key_b64)
    if len(key) != 32:
        raise SystemExit(f"密钥长度错误: {len(key)} bytes,需要 32 bytes (AES-256)")
    return key


def encrypt_field(aesgcm: AESGCM, value: str) -> str:
    """加密单个文本字段,返回 "enc:v1:<b64 nonce>:<b64 ciphertext>"。"""
    nonce = os.urandom(NONCE_LEN)
    ct = aesgcm.encrypt(nonce, value.encode("utf-8"), None)
    return "enc:v1:" + base64.b64encode(nonce).decode() + ":" + base64.b64encode(ct).decode()


def is_encrypted(value: str) -> bool:
    return value.startswith("enc:v1:")


def encrypt_file(aesgcm: AESGCM, data: bytes) -> bytes:
    """整文件 AES-GCM:enc:v1 头 + nonce + ciphertext(二进制输出)。"""
    nonce = os.urandom(NONCE_LEN)
    ct = aesgcm.encrypt(nonce, data, None)
    return b"enc:v1:" + base64.b64encode(nonce) + b":" + base64.b64encode(ct)


def decrypt_file(aesgcm: AESGCM, blob: bytes) -> bytes:
    assert blob.startswith(b"enc:v1:")
    _, nonce_b64, ct_b64 = blob.split(b":")
    return aesgcm.decrypt(base64.b64decode(nonce_b64), base64.b64decode(ct_b64), None)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--key-file", required=True, type=Path)
    parser.add_argument("--fields", default="properties_json,assertion_text,evidence_text,text")
    parser.add_argument("--file-encrypt", action="store_true",
                        help="整文件 AES-GCM 加密(parquet 等二进制,无法逐字段处理)")
    args = parser.parse_args()

    key = load_key(args.key_file)
    aesgcm = AESGCM(key)

    if args.file_encrypt:
        data = args.input.read_bytes()
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_bytes(encrypt_file(aesgcm, data))
        print(f"整文件加密完成: {args.input} ({len(data)} bytes) -> {args.output}")
        return 0

    key = load_key(args.key_file)
    aesgcm = AESGCM(key)
    fields = [f.strip() for f in args.fields.split(",") if f.strip()]

    encrypted = 0
    total = 0
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.input.open("r", encoding="utf-8") as fin, args.output.open("w", encoding="utf-8") as fout:
        for line in fin:
            total += 1
            line = line.rstrip("\n")
            if not line.strip():
                fout.write("\n")
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                fout.write(line + "\n")
                continue
            for f in fields:
                if f in rec and isinstance(rec[f], str) and rec[f] and not is_encrypted(rec[f]):
                    rec[f] = encrypt_field(aesgcm, rec[f])
                    encrypted += 1
            fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
    print(f"KB 加密回填完成: {total} 行, {encrypted} 个字段加密")
    return 0


if __name__ == "__main__":
    sys.exit(main())