#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""newapi_pool_provision.py -- 内测 new-api per-user ¥100 额度池离线预置脚本（WU2 · 接线点 B）

权威设计：docs/agent-specs/2026-07-07-内测-WU2-newapi额度接入-设计.md §3.3（离线 ops 预置脚本）
        + §3.5（池表 newapi_quota_pool 契约）。

它做什么
--------
在 mini-infra 上一次性预建 N 个专属 (new-api user + token + ¥100 额度) 账户，
把每个条目 emit 成可导入 game-cloud 池表 newapi_quota_pool 的 JSON + SQL，供 WU2 落地后灌数。
provision 全部离线做（S0 实测坐实：new-api admin 令牌不能替他人建 token、access_token
不可经 API 设，per-user token 必须直连其 postgres），运行时零 new-api admin 耦合。

每个池条目四步建成（S0 坐实的唯一可行路径）
------------------------------------------------
  1. POST /api/user/（root 令牌）建用户，username = 确定性 `<前缀><序号>`；从 postgres 查 uid。
  2. postgres 直写 UPDATE users SET access_token=<确定性 32 位串> WHERE id=uid
     —— access_token 无法经 API 设，又是下一步「以该用户身份建 token」的前提。
  3. POST /api/token/（以该用户 access_token + New-Api-User:uid 建 token），
     remain_quota=¥100 折算、unlimited_quota=false、expired_time=-1（不过期）；从 postgres 查 token_id/key。
  4. PUT /api/user/（root 令牌）设 user.quota=¥100 折算（账户自洽；权威余额闸仍是 token.remain_quota）。

幂等（可重跑补池）
------------------
按确定性 username 去重：重跑时若用户已存在则复用其 uid，逐步「补齐」缺失的 access_token /
token / quota（ensure 语义），不重复建号。补池 = 提高 --count 重跑。单条失败只记该条、不污染后续。

外部交互红线
------------
  - HTTP：连接/读超时（urlopen timeout=读10s，内部连粒度由 socket 兜），非 2xx / success=false 归错，
    5xx/超时重试 2 次（指数退避），4xx 不重试；全程绕系统代理（ProxyHandler({})，内网直连 100.64.x）。
  - postgres：经 `docker exec infra-postgres psql` 执行（S0 验证路径），失败即抛、该条标记失败。
  - 日志：token / access_token 一律脱敏（前后各留 4 位）；可追溯（时间戳 + 步骤 + 条目序号）。

凭据（不硬编码）
----------------
  - postgres 密码：环境变量 NEWAPI_PG_PASSWORD（必填；见 docs/内网凭据与端点.md）。
  - root 管理令牌：环境变量 NEWAPI_ROOT_TOKEN（选填）；未设则自动从 postgres users.id=1 读。

用法（在 mini-infra 上跑，docker 本地可用）
------------------------------------------
  export NEWAPI_PG_PASSWORD=<见凭据文档>
  # 预置 2 条测试条目（验证用，前缀明确可辨），dry-run 先看：
  python3 newapi_pool_provision.py --count 2 --prefix neice_s0test_ --dry-run
  python3 newapi_pool_provision.py --count 2 --prefix neice_s0test_
  # 清理测试条目（删 tokens + users 行）：
  python3 newapi_pool_provision.py --cleanup 'neice_s0test_%'
"""
from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import re
import secrets
import subprocess
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path

# ─── 常量 ────────────────────────────────────────────────────────────────
NEWAPI_BASE = os.environ.get("NEWAPI_BASE", "http://localhost:3000")  # 脚本宿主 mini-infra，本地直连
PG_CONTAINER = os.environ.get("NEWAPI_PG_CONTAINER", "infra-postgres")
PG_DB = os.environ.get("NEWAPI_PG_DB", "new-api")
PG_USER = os.environ.get("NEWAPI_PG_USER", "root")
ROOT_UID = 1  # new-api root 用户 id（qingse），持系统管理 access_token
HTTP_CONNECT_TIMEOUT = 5.0
HTTP_READ_TIMEOUT = 10.0
HTTP_RETRIES = 2  # 5xx / 超时最多重试 2 次
USERNAME_RE = re.compile(r"^[A-Za-z0-9_]+$")  # 只允许安全 username，杜绝 SQL/命令注入
POOL_TABLE = "newapi_quota_pool"

log = logging.getLogger("newapi_pool")


def _mask(token: str | None) -> str:
    """脱敏敏感串：前后各留 4 位，中间打码。用于日志（红线：token/access_token 不整条打日志）。"""
    if not token:
        return "<empty>"
    if len(token) <= 10:
        return token[:2] + "***"
    return f"{token[:4]}…{token[-4:]}(len={len(token)})"


# ─── postgres 直连（经 docker exec psql，S0 验证路径） ──────────────────────
class Postgres:
    """封装对 new-api postgres 的只读/写操作。所有写走事务语义（单条 SQL 天然原子）。"""

    def __init__(self, password: str):
        if not password:
            raise RuntimeError(
                "缺少 NEWAPI_PG_PASSWORD 环境变量——无法连接 new-api postgres。"
                "见 docs/内网凭据与端点.md new-api 段。"
            )
        self._password = password

    def _run(self, sql: str) -> str:
        """执行一条 SQL，返回 tuples-only、unaligned、以 | 分隔的原始 stdout。失败即抛。"""
        cmd = [
            "docker", "exec",
            "-e", f"PGPASSWORD={self._password}",
            PG_CONTAINER,
            "psql", "-U", PG_USER, "-d", PG_DB,
            "-t", "-A", "-F", "|", "-c", sql,
        ]
        try:
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        except subprocess.TimeoutExpired as e:
            raise RuntimeError(f"postgres 执行超时：{sql[:80]}") from e
        if proc.returncode != 0:
            # 可追溯错误日志：暴露 psql stderr，但不含敏感值（SQL 里的 access_token 已在调用侧脱敏日志）
            raise RuntimeError(f"postgres 执行失败 rc={proc.returncode}: {proc.stderr.strip()}")
        return proc.stdout.strip()

    def query_root_token(self) -> str:
        out = self._run(f"SELECT access_token FROM users WHERE id={ROOT_UID};")
        if not out:
            raise RuntimeError("postgres 未查到 root(id=1) access_token")
        return out.splitlines()[0].strip()

    def query_user_id(self, username: str) -> int | None:
        out = self._run(f"SELECT id FROM users WHERE username='{username}';")
        line = out.splitlines()[0].strip() if out else ""
        return int(line) if line else None

    def query_access_token(self, uid: int) -> str:
        out = self._run(f"SELECT access_token FROM users WHERE id={uid};")
        return out.splitlines()[0].strip() if out else ""

    def query_user_quota(self, uid: int) -> int | None:
        out = self._run(f"SELECT quota FROM users WHERE id={uid};")
        line = out.splitlines()[0].strip() if out else ""
        return int(line) if line else None

    def set_access_token(self, uid: int, access_token: str) -> None:
        # access_token 无法经 API 设，必须 DB 直写（S0 坐实的关键步）
        self._run(f"UPDATE users SET access_token='{access_token}' WHERE id={uid};")

    def query_token(self, uid: int, name: str) -> tuple[int, str] | None:
        """按 user_id + token name 查 token，返回 (token_id, token_key)。"""
        out = self._run(
            f"SELECT id, key FROM tokens WHERE user_id={uid} AND name='{name}' ORDER BY id LIMIT 1;"
        )
        line = out.splitlines()[0].strip() if out else ""
        if not line or "|" not in line:
            return None
        tid, key = line.split("|", 1)
        return int(tid), key.strip()

    def delete_user_cascade(self, username_like: str) -> tuple[int, int]:
        """清理：按 username LIKE 删 users + 其 tokens。返回 (删 tokens 数, 删 users 数)。
        安全护栏：拒绝任何可能命中 root(id<=1) 的模式；只允许显式 % 通配。"""
        # 先查将被删的 uid 列表，护栏校验
        uids_out = self._run(f"SELECT id FROM users WHERE username LIKE '{username_like}';")
        uids = [int(x) for x in uids_out.splitlines() if x.strip()]
        if not uids:
            return (0, 0)
        if any(u <= ROOT_UID for u in uids):
            raise RuntimeError(f"清理护栏：模式 {username_like} 命中受保护用户 id<={ROOT_UID}，拒绝执行")
        uid_csv = ",".join(str(u) for u in uids)
        tok_out = self._run(f"WITH d AS (DELETE FROM tokens WHERE user_id IN ({uid_csv}) RETURNING 1) SELECT count(*) FROM d;")
        usr_out = self._run(f"WITH d AS (DELETE FROM users WHERE id IN ({uid_csv}) RETURNING 1) SELECT count(*) FROM d;")
        return (int(tok_out.strip() or 0), int(usr_out.strip() or 0))


# ─── new-api HTTP（stdlib urllib，绕系统代理，带超时+重试） ──────────────────
_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))  # 空 ProxyHandler = 禁用系统代理


def _http_json(method: str, path: str, token: str, newapi_user: int, body: dict | None = None) -> dict:
    """对 new-api 发一个 JSON 请求。红线：绕代理 / 超时 / 非 2xx 归错 / 5xx 重试。

    path 必须带尾斜杠（new-api gin RedirectTrailingSlash，否则 307）。
    鉴权头：Authorization: Bearer <token> + New-Api-User: <uid>（缺后者即 401）。
    """
    url = NEWAPI_BASE.rstrip("/") + path
    data = json.dumps(body).encode("utf-8") if body is not None else None
    headers = {
        "Authorization": f"Bearer {token}",
        "New-Api-User": str(newapi_user),
        "Content-Type": "application/json",
    }
    last_err: Exception | None = None
    for attempt in range(HTTP_RETRIES + 1):
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with _OPENER.open(req, timeout=HTTP_READ_TIMEOUT) as resp:
                raw = resp.read().decode("utf-8")
                payload = json.loads(raw) if raw else {}
                # new-api 标准返回体 {success,message,data}；success=false 视为业务错
                if isinstance(payload, dict) and payload.get("success") is False:
                    raise RuntimeError(f"{method} {path} 业务失败: {payload.get('message')}")
                return payload
        except urllib.error.HTTPError as e:
            body_txt = e.read().decode("utf-8", "ignore")[:200]
            last_err = RuntimeError(f"{method} {path} HTTP {e.code}: {body_txt}")
            if 500 <= e.code < 600 and attempt < HTTP_RETRIES:
                time.sleep(0.5 * (2 ** attempt))
                log.warning("[http] %s %s 5xx 重试 %d/%d", method, path, attempt + 1, HTTP_RETRIES)
                continue
            raise last_err  # 4xx 不重试
        except (urllib.error.URLError, TimeoutError, OSError) as e:
            last_err = RuntimeError(f"{method} {path} 网络错误: {e}")
            if attempt < HTTP_RETRIES:
                time.sleep(0.5 * (2 ** attempt))
                log.warning("[http] %s %s 网络错重试 %d/%d: %s", method, path, attempt + 1, HTTP_RETRIES, e)
                continue
            raise last_err
    raise last_err or RuntimeError("unreachable")


# ─── ¥ ↔ quota 折算（口径与 cost.py 同源，读网关而非硬编码） ─────────────────
def fetch_conversion() -> tuple[int, float]:
    """GET /api/status 拿 quota_per_unit / usd_exchange_rate（无需鉴权）。缺 usd 用 7.3 默认。"""
    url = NEWAPI_BASE.rstrip("/") + "/api/status"
    req = urllib.request.Request(url, method="GET")
    with _OPENER.open(req, timeout=HTTP_READ_TIMEOUT) as resp:
        data = json.loads(resp.read().decode("utf-8")).get("data", {})
    qpu = int(data.get("quota_per_unit") or 500000)
    usd = data.get("usd_exchange_rate")
    usd = float(usd) if usd else 7.3  # DB options 未覆盖时默认 7.3
    return qpu, usd


def compute_grant(yuan: float, qpu: int, usd: float) -> int:
    """¥ → quota：round(yuan / usd × qpu)。¥100 @ 500000/7.3 ≈ 6,849,315。"""
    return round(yuan / usd * qpu)


# ─── 单条目预置（ensure 语义，幂等可补齐） ──────────────────────────────────
def provision_one(pg: Postgres, root_token: str, username: str, grant: int,
                  qpu: int, usd: float, dry_run: bool) -> dict:
    """预置/补齐一个池条目，返回条目 dict（对齐 §3.5 池表契约）。任一步失败即抛，由调用方记为失败。"""
    if not USERNAME_RE.match(username):
        raise RuntimeError(f"非法 username: {username}")

    if dry_run:
        log.info("[dry-run] 计划预置 %s（grant=%d, unlimited=false, expired=-1）", username, grant)
        return {"username": username, "dry_run": True, "grant_quota": grant}

    # 步骤 1：建用户（幂等——已存在则复用）
    uid = pg.query_user_id(username)
    if uid is None:
        password = secrets.token_urlsafe(12)  # 随机密码 16 位（new-api Password max=20）；仅占位，从不密码登录
        _http_json("POST", "/api/user/", root_token, ROOT_UID, {
            "username": username,
            "password": password,
            "display_name": username[:20],  # new-api DisplayName max=20，直接用 username（截断兜底）
        })
        uid = pg.query_user_id(username)
        if uid is None:
            raise RuntimeError(f"建用户 {username} 后 postgres 未查到 uid")
        log.info("[1/4] 建用户 %s → uid=%d", username, uid)
    else:
        log.info("[1/4] 用户 %s 已存在 uid=%d（复用）", username, uid)

    # 步骤 2：DB 直写 access_token（确定性，可重复 UPDATE 幂等）
    access_token = hashlib.sha256(f"neice-quota-pool::v1::{username}".encode()).hexdigest()[:32]
    existing_at = pg.query_access_token(uid)
    if existing_at != access_token:
        pg.set_access_token(uid, access_token)
        log.info("[2/4] DB 写 access_token uid=%d → %s", uid, _mask(access_token))
    else:
        log.info("[2/4] access_token uid=%d 已就绪 %s（复用）", uid, _mask(access_token))

    # 步骤 3：以该用户身份建 token（幂等——已存在同名 token 则复用）
    tok = pg.query_token(uid, username)
    if tok is None:
        _http_json("POST", "/api/token/", access_token, uid, {
            "name": username,
            "remain_quota": grant,
            "unlimited_quota": False,
            "expired_time": -1,  # 不过期，避免自伤式过期（设计 §3.9）
        })
        tok = pg.query_token(uid, username)
        if tok is None:
            raise RuntimeError(f"建 token 后 postgres 未查到 user={uid} name={username}")
        log.info("[3/4] 建 token uid=%d → token_id=%d key=%s", uid, tok[0], _mask(tok[1]))
    else:
        log.info("[3/4] token uid=%d name=%s 已存在 token_id=%d（复用）", uid, username, tok[0])
    token_id, token_key = tok

    # 步骤 4：设 user.quota（root 令牌；账户自洽，权威余额闸仍是 token.remain_quota）
    cur_quota = pg.query_user_quota(uid)
    if cur_quota != grant:
        _http_json("PUT", "/api/user/", root_token, ROOT_UID, {
            "id": uid,
            "username": username,
            "display_name": username[:20],
            "quota": grant,
            "group": "default",
        })
        log.info("[4/4] 设 user.quota uid=%d → %d", uid, grant)
    else:
        log.info("[4/4] user.quota uid=%d 已=%d（复用）", uid, grant)

    # 组装池表条目（§3.5 契约字段）
    return {
        "newapi_user_id": uid,
        "newapi_token_id": token_id,
        "newapi_token_key": token_key,
        "grant_quota": grant,
        "quota_per_unit_snapshot": qpu,
        "usd_rate_snapshot": usd,
        "status": "FREE",
        "username": username,  # 仅审计参考，非池表列
    }


# ─── emit：JSON + SQL（对齐 §3.5，供 WU2 落地后导入 game-cloud） ─────────────
def emit_outputs(entries: list[dict], out_dir: Path, prefix: str) -> tuple[Path, Path]:
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    out_dir.mkdir(parents=True, exist_ok=True)
    stem = f"newapi_pool_{prefix.rstrip('_')}_{ts}"
    json_path = out_dir / f"{stem}.json"
    sql_path = out_dir / f"{stem}.sql"

    # 只保留池表列（去掉 username 审计字段）
    pool_cols = ["newapi_user_id", "newapi_token_id", "newapi_token_key",
                 "grant_quota", "quota_per_unit_snapshot", "usd_rate_snapshot", "status"]
    json_rows = [{k: e[k] for k in pool_cols} for e in entries]
    json_path.write_text(json.dumps(json_rows, ensure_ascii=False, indent=2), encoding="utf-8")

    # SQL：INSERT ... ON DUPLICATE KEY UPDATE（去重键 uk(newapi_user_id)，§3.3(a)）
    # 刻意不在 ON DUPLICATE 里改 status/claimed_*——重跑不把已 CLAIMED 条目刷回 FREE（防免费续杯 / 防误 unclaim）。
    lines = [
        f"-- 内测 new-api 额度池预置产物 {ts}，共 {len(json_rows)} 条。",
        f"-- 对齐 docs/agent-specs/2026-07-07-内测-WU2-newapi额度接入-设计.md §3.5 池表 {POOL_TABLE}（V32 待落）。",
        f"-- 导入前提：game-cloud 已建 {POOL_TABLE} 表（含 uk(newapi_user_id)、审计列默认值）。",
        f"-- ON DUPLICATE 只更 token/额度快照，不动 status/claimed_*（幂等补池不 unclaim 已绑条目）。",
    ]
    for r in json_rows:
        key_sql = r["newapi_token_key"].replace("'", "''")
        lines.append(
            f"INSERT INTO {POOL_TABLE} "
            f"(newapi_user_id, newapi_token_id, newapi_token_key, grant_quota, "
            f"quota_per_unit_snapshot, usd_rate_snapshot, status) VALUES "
            f"({r['newapi_user_id']}, {r['newapi_token_id']}, '{key_sql}', {r['grant_quota']}, "
            f"{r['quota_per_unit_snapshot']}, {r['usd_rate_snapshot']}, '{r['status']}') "
            f"ON DUPLICATE KEY UPDATE "
            f"newapi_token_id=VALUES(newapi_token_id), newapi_token_key=VALUES(newapi_token_key), "
            f"grant_quota=VALUES(grant_quota), quota_per_unit_snapshot=VALUES(quota_per_unit_snapshot), "
            f"usd_rate_snapshot=VALUES(usd_rate_snapshot);"
        )
    sql_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return json_path, sql_path


# ─── 主流程 ──────────────────────────────────────────────────────────────
def cmd_provision(args) -> int:
    pg = Postgres(os.environ.get("NEWAPI_PG_PASSWORD", ""))
    root_token = os.environ.get("NEWAPI_ROOT_TOKEN") or pg.query_root_token()
    log.info("root 管理令牌就绪 %s", _mask(root_token))

    qpu, usd = fetch_conversion()
    grant = compute_grant(args.yuan, qpu, usd)
    log.info("折算：¥%.0f @ quota_per_unit=%d usd=%.4f → grant_quota=%d", args.yuan, qpu, usd, grant)

    entries: list[dict] = []
    failures: list[dict] = []
    for i in range(args.start, args.start + args.count):
        username = f"{args.prefix}{i:03d}"
        try:
            entries.append(provision_one(pg, root_token, username, grant, qpu, usd, args.dry_run))
        except Exception as e:  # 单条失败不污染后续（红线）
            log.error("[条目 %s] 预置失败：%s", username, e)
            failures.append({"username": username, "error": str(e)})

    log.info("完成：成功 %d / 失败 %d", len(entries), len(failures))
    if args.dry_run:
        log.info("dry-run 结束，未写 new-api、未 emit 文件。")
        return 0 if not failures else 1

    if entries:
        json_path, sql_path = emit_outputs(entries, Path(args.out_dir), args.prefix)
        log.info("emit JSON → %s", json_path)
        log.info("emit SQL  → %s", sql_path)
    if failures:
        log.warning("失败条目：%s", json.dumps(failures, ensure_ascii=False))
        return 1
    return 0


def cmd_cleanup(args) -> int:
    """清理：按 username LIKE 模式删 tokens + users（护栏拒绝命中 root）。"""
    pg = Postgres(os.environ.get("NEWAPI_PG_PASSWORD", ""))
    ntok, nusr = pg.delete_user_cascade(args.cleanup)
    log.info("清理模式 %s：删 tokens %d 行 / users %d 行", args.cleanup, ntok, nusr)
    return 0


def main() -> int:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        stream=sys.stderr,
    )
    p = argparse.ArgumentParser(description="内测 new-api per-user ¥100 额度池离线预置")
    p.add_argument("--count", type=int, default=1, help="目标条目数 N（本次预置到 start..start+count）")
    p.add_argument("--start", type=int, default=1, help="起始序号（默认 1）")
    p.add_argument("--prefix", default="neice_", help="username 前缀（默认 neice_；测试用 neice_s0test_）")
    p.add_argument("--yuan", type=float, default=100.0, help="每条目额度（元，默认 100）")
    p.add_argument("--out-dir", default="./newapi-pool-out", help="JSON/SQL 产物目录")
    p.add_argument("--dry-run", action="store_true", help="只打印计划，不写 new-api、不 emit")
    p.add_argument("--cleanup", metavar="LIKE_PATTERN",
                   help="清理模式：删匹配 username LIKE 的 users+tokens（如 'neice_s0test_%%'）")
    args = p.parse_args()

    if args.cleanup:
        return cmd_cleanup(args)
    return cmd_provision(args)


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