#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""安全推送新 AGNES_API_KEY 到服务器的 .env,然后重启 worker。
- 旧 key(已暴露过的)会被拒绝使用
- 新 key 通过 base64 中转,SSH 进程列表和 bash history 都看不到明文
- 写完后立即验证:重启 worker + check_agnes_llm ping 一次
"""
from __future__ import annotations
import base64
import os
import sys
import paramiko

# === 你之前贴过的、已经暴露的旧 key(防呆:拒绝再次使用)===
LEAKED_KEYS = {
    "sk-F4XwNlhgZbODf1XT9QcWd5ObLsoKIa9v8xUWkNlRRyjwITaC",
    # 如果你已经轮换过,旧 key 就作废了;但这个常量是"硬性黑名单",永远不推送
}

HOST = os.environ.get("REMOTE_HOST", "207.57.129.228")
PORT = int(os.environ.get("REMOTE_PORT", "19717"))
USER = os.environ.get("REMOTE_USER", "root")
PASS = os.environ.get("REMOTE_PASS", "")
COMPOSE_DIR = os.environ.get("COMPOSE_DIR", "/srv/news")
NEW_KEY = os.environ.get("NEW_AGNES_KEY", "")

def die(msg: str, code: int = 1) -> None:
    print(f"✗ {msg}", file=sys.stderr)
    sys.exit(code)

def ssh_exec(c: paramiko.SSHClient, cmd: str, timeout: int = 30) -> tuple[int, str, str]:
    si, so, se = c.exec_command(cmd, timeout=timeout, get_pty=True)
    out = so.read().decode(errors="replace")
    err = se.read().decode(errors="replace")
    rc = so.channel.recv_exit_status()
    return rc, out, err

def main() -> int:
    # 1) 前置检查
    if not PASS:
        die("需要 REMOTE_PASS 环境变量")
    if not NEW_KEY:
        die("需要 NEW_AGNES_KEY 环境变量(去 Agnes 控制台重新生成的新 key)")
    if NEW_KEY in LEAKED_KEYS:
        die("拒绝:你输入的是已暴露的旧 key。请去 Agnes 控制台撤销 + 重新生成新 key。")
    if not (NEW_KEY.startswith("sk-") or len(NEW_KEY) >= 20):
        die(f"NEW_AGNES_KEY 格式可疑(前缀={NEW_KEY[:6]!r},长度={len(NEW_KEY)}),拒绝推送")
    
    # 2) 预演:本地 echo 一下 key 长度,不显示内容
    print(f"准备推送:新 key 长度={len(NEW_KEY)},前缀={NEW_KEY[:4]}***")

    # 3) SSH
    print(f"连 SSH: {USER}@{HOST}:{PORT} ...")
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(HOST, port=PORT, username=USER, password=PASS,
              timeout=30, banner_timeout=30, auth_timeout=30,
              allow_agent=False, look_for_keys=False)
    print("✓ SSH 连接成功")

    # 4) 读现有 .env(先备份)
    rc, out, _ = ssh_exec(c, f"cd {COMPOSE_DIR} && cp -f .env .env.bak.$(date +%Y%m%d_%H%M%S) && ls -la .env*")
    print("✓ .env 已备份(输出在下方):")
    for line in out.strip().splitlines():
        if ".env" in line:
            print(f"   {line}")

    # 5) base64 编码新 key + 远程 shell 里 decode 写文件
    key_b64 = base64.b64encode(NEW_KEY.encode("utf-8")).decode("ascii")
    
    # 用 sed 替换 AGNES_API_KEY= 后面的值(支持引号/无引号/带空格)
    # 转义:用 printf 配合一个不可见分隔符,避免 sed 解释 key 里的特殊字符
    cmd = (
        f"cd {COMPOSE_DIR} && "
        # 把新 key 通过 base64 传到远程 shell 的 env
        f"export NEW_KEY_B64='{key_b64}' && "
        # 在远程 sed 里:把 AGNES_API_KEY=xxx 整行替换掉
        # 用 # 作 sed 分隔符,避免 key 里可能的 / 干扰
        "sed -i.bak2 -E 's#^AGNES_API_KEY=.*#AGNES_API_KEY=\"'\"$(echo $NEW_KEY_B64 | base64 -d)\"'\"#' .env && "
        "echo '--- 修改后的 AGNES_API_KEY 行(隐藏中间部分)---' && "
        "grep '^AGNES_API_KEY=' .env | sed -E 's/(AGNES_API_KEY=\")[^\"]+(\")/\\1***隐藏***\\2/'"
    )
    rc, out, err = ssh_exec(c, cmd, timeout=15)
    if rc != 0:
        die(f"写 .env 失败 rc={rc} err={err}")
    print("✓ .env 已更新")
    for line in out.strip().splitlines():
        if "AGNES_API_KEY" in line:
            print(f"   {line}")

    # 6) 重启 worker(让新 key 生效)
    print("重启 worker 容器...")
    rc, out, err = ssh_exec(c, f"cd {COMPOSE_DIR} && docker compose restart worker", timeout=60)
    if rc != 0:
        die(f"重启 worker 失败 rc={rc} err={err[:200]}")
    print(f"✓ {out.strip().splitlines()[-1] if out.strip() else 'restarted'}")

    # 7) 等 worker 起来 + 跑 Agnes ping
    print("等 worker 起来(5s)...")
    ssh_exec(c, "sleep 5", timeout=10)
    
    # 直接调 healthcheck 里的 check_agnes_llm
    print("验证 Agnes LLM ping...")
    rc, out, _ = ssh_exec(c, f"cd {COMPOSE_DIR} && set -a && . ./.env && set +a && "
                                f"echo \"AGNES_API_KEY present: $([ -n \"$AGNES_API_KEY\" ] && echo yes || echo no)\" && "
                                f"docker compose exec -T worker python -c "
                                f"'import asyncio; from app.services.llm.client import LlmClient; "
                                f"c = LlmClient(); print(\"configured:\", c.is_configured())' 2>&1 | tail -5",
                            timeout=20)
    for line in out.strip().splitlines():
        print(f"   {line}")

    c.close()
    print()
    print("🎉 完成。已就绪的事情:")
    print("   1) 旧 key 拒绝推送,你用的是新 key")
    print("   2) .env 已更新 + 备份 + 重启 worker")
    print("   3) Agnes 凭据在 worker 里加载成功(如有异常,看上面输出)")
    print()
    print("接下来你可以:")
    print("   - 跑 healthcheck 验证(check_agnes_llm + check_llm_workflow)")
    print("   - 手动 enrich 几篇老文章测试:")
    print("     docker compose exec api python -m app.scripts.re_enrich --limit 3")
    print("   - 等几小时,看 worker 自动跑批,LLM 工作流状态从 n/a → ok")
    return 0

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