"""MASTER_KEY 丢失恢复（一次性脚本，日常运行不接触）

用法：
  RESTORE_KEY=<钥匙> .venv/bin/python scripts/recover_master_key.py           # 仅打印找回的 MASTER_KEY
  RESTORE_KEY=<钥匙> .venv/bin/python scripts/recover_master_key.py --write   # 备份 .env 后回写 MASTER_KEY

也支持 --key 传参（会残留 shell 历史，建议用环境变量）。
恢复后重启服务即可：全部密文零迁移、照常解密。
"""

import argparse
import os
import re
import shutil
import sys
from pathlib import Path

# 确保能导入 app 包（脚本位于 scripts/ 子目录）
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.core import crypto

BASE_DIR = Path(__file__).resolve().parent.parent
DEFAULT_ESCROW = BASE_DIR / "data" / "master_key.escrow"
DEFAULT_ENV = BASE_DIR / ".env"


def _write_env(master_key: str, env_path: Path) -> Path | None:
    """备份 .env 后回写 MASTER_KEY 行（不存在则追加），返回备份路径"""
    backup: Path | None = None
    if env_path.exists():
        backup = env_path.parent / (env_path.name + ".bak-pre-recover")
        shutil.copy2(env_path, backup)
        text = env_path.read_text(encoding="utf-8")
        if re.search(r"^MASTER_KEY=.*$", text, flags=re.MULTILINE):
            text = re.sub(
                r"^MASTER_KEY=.*$", f"MASTER_KEY={master_key}", text, count=1, flags=re.MULTILINE
            )
        else:
            text = text.rstrip("\n") + f"\nMASTER_KEY={master_key}\n"
        env_path.write_text(text, encoding="utf-8")
    else:
        env_path.write_text(f"MASTER_KEY={master_key}\n", encoding="utf-8")
    return backup


def main() -> int:
    parser = argparse.ArgumentParser(description="凭 RESTORE_KEY 找回 MASTER_KEY")
    parser.add_argument(
        "--key", help="RESTORE_KEY（建议改用环境变量 RESTORE_KEY，避免 shell 历史残留）"
    )
    parser.add_argument(
        "--write", action="store_true", help="找回后直接回写 .env（先备份为 .env.bak-pre-recover）"
    )
    parser.add_argument(
        "--escrow", default=str(DEFAULT_ESCROW),
        help="escrow 文件路径（默认 data/master_key.escrow）",
    )
    parser.add_argument("--env", default=str(DEFAULT_ENV), help=".env 路径（默认项目根 .env）")
    args = parser.parse_args()

    restore_key = args.key or os.environ.get("RESTORE_KEY")
    if not restore_key:
        print("[恢复] 缺少 RESTORE_KEY：--key 传入或设置环境变量 RESTORE_KEY", file=sys.stderr)
        return 2
    escrow_path = Path(args.escrow)
    if not escrow_path.exists():
        print(f"[恢复] 未找到托管文件 {escrow_path}（尚未建立托管？）", file=sys.stderr)
        return 2

    try:
        master_key = crypto.recover_master_key(
            restore_key.strip(), escrow_path.read_text(encoding="utf-8").strip()
        )
    except ValueError as e:
        print(f"[恢复] 失败：{e}", file=sys.stderr)
        return 1

    if args.write:
        backup = _write_env(master_key, Path(args.env))
        where = f"（原文件已备份至 {backup.name}）" if backup else ""
        print(f"[恢复] MASTER_KEY 已找回并回写 {args.env}{where}")
        print("[恢复] 请重启服务使 key 生效：sudo systemctl restart vps-manager")
    else:
        print("[恢复] MASTER_KEY 已找回（未回写 .env）：")
        print(f"  MASTER_KEY={master_key}")
        print("[恢复] 确认无误后加 --write 自动回写 .env，或手动粘贴到 .env 的 MASTER_KEY= 行")
    return 0


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