#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
门禁核验（POST /api/v1/mp/roster/match）验收测试 —— M3 的第一块。

覆盖《03-接口设计》§5.4：
    [x] 姓名 + 证件号完全匹配 → 放行并绑定
    [x] 返回里**只有脱敏证件号**，绝无明文
    [x] 幂等：已绑自己的人再核验仍成功（学生刷新页面不该看到吓人的提示）
    [x] **失败提示统一**：姓名错与号码错给**完全相同**的一句话（防身份证号枚举）
    [x] 40072：同一名册项被另一个账号核验 → 已被登记
    [x] **并发**：两个账号同时对同一行核验 → 恰好一个成功
    [x] 限流：同账号 10 次/小时 → 40071
    [x] 跨租户：拿别租户的 taskId 不泄露
    [x] 预填：老师已填的报考职业/工种/等级随核验结果返回

怎么拿 token（本脚本的关键手法）：
    门禁要求"已登录的学生"，而登录（/mp/login）需要真实的 wx.login code —— 脚本拿不到。
    所以这里**用 dev 配置里的 JWT 密钥自己签一个 MP token**（HS512，claims 与
    JwtService.issue 一致）。这与 e2e-m0 用 dev RSA 公钥加密登录密码是同一类做法：
    都是"用开发配置里的凭据手搓一个测试夹具"，不是绕过认证逻辑。
    ⚠️ 它只对 dev 有效；prod 的密钥不在仓库里，这个手法自然失效。

用法
    python scripts/e2e-m3-gate.py

依赖：cryptography；且需要能 `docker exec xt-pg` / `xt-redis`
退出码：0 = 全部通过
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import json
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request
import uuid

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PUB_KEY = os.path.join(ROOT, "deploy", "rsa-public-key-dev.pem")
DEV_YML = os.path.join(ROOT, "backend", "xt-server", "src", "main", "resources", "application-dev.yml")

PG_CONTAINER = os.environ.get("XT_PG_CONTAINER", "xt-pg")
PG_USER = os.environ.get("PG_USER", "fill")
PG_DB = os.environ.get("PG_DB", "filldb")
PG_PWD = os.environ.get("PG_PASSWORD", "fill_dev_pwd")
REDIS_CONTAINER = os.environ.get("XT_REDIS_CONTAINER", "xt-redis")
REDIS_PWD = os.environ.get("REDIS_PASSWORD", "redis_dev_pwd")

PASSED: list[str] = []
FAILED: list[tuple[str, str]] = []

CREATED_TASK_IDS: list[int] = []
CREATED_USER_IDS: list[int] = []

# 名册里的两条测试数据（证件号都校验位正确）
ID_A = "361029200912100015"
NAME_A = "曾伊豪"
ID_B = "440524199001010018"
NAME_B = "杨晨"


def ok(name: str) -> None:
    PASSED.append(name)
    print(f"  [OK]   {name}")


def fail(name: str, detail: str) -> None:
    FAILED.append((name, detail))
    print(f"  [FAIL] {name}\n         {detail}")


def check(name: str, cond: bool, detail: str = "") -> bool:
    if cond:
        ok(name)
    else:
        fail(name, detail or "断言不成立")
    return cond


def sql(stmt: str) -> tuple[bool, str]:
    try:
        r = subprocess.run(
            ["docker", "exec", "-e", f"PGPASSWORD={PG_PWD}", PG_CONTAINER,
             "psql", "-U", PG_USER, "-d", PG_DB, "-t", "-A", "-c", stmt],
            capture_output=True, text=True, timeout=40)
        return r.returncode == 0, (r.stdout or r.stderr).strip()
    except Exception as e:
        return False, str(e)


def redis(args: list[str]) -> tuple[bool, str]:
    try:
        r = subprocess.run(
            ["docker", "exec", "-e", f"REDISCLI_AUTH={REDIS_PWD}", REDIS_CONTAINER,
             "redis-cli"] + args, capture_output=True, text=True, timeout=30)
        return r.returncode == 0, (r.stdout or r.stderr).strip()
    except Exception as e:
        return False, str(e)


def clear_rate() -> None:
    """清掉门禁限流键 —— 否则本脚本自己会被限流（它要发十几次请求）。"""
    okk, out = redis(["--scan", "--pattern", "xt:mp:gate:*"])
    if okk and out.strip():
        for k in out.strip().splitlines():
            redis(["DEL", k])


def cleanup_all() -> None:
    if CREATED_TASK_IDS:
        ids = ",".join(str(i) for i in CREATED_TASK_IDS)
        for stmt in [
            f"DELETE FROM collect_roster WHERE task_id IN ({ids})",
            f"DELETE FROM collect_task WHERE id IN ({ids})",
        ]:
            sql(stmt)
    base = CTX.get("base")
    token = CTX.get("token")
    if base and token:
        for uid in CREATED_USER_IDS:
            try:
                call("DELETE", f"{base}/api/v1/users/{uid}", token)
            except Exception:
                pass
    clear_rate()


CTX: dict[str, str | None] = {"base": None, "token": None}


# ── 配置读取与 token 手搓 ────────────────────────────

def read_dev_config() -> dict:
    """从 application-dev.yml（+ application.yml 兜底）里读 jwt.secret / hmac-key / issuer。

    ⚠️ 要读<b>两个</b>文件：`secret` 与 `hmac-key` 在 dev 配置里，
    而 `issuer` 在主配置 application.yml 里（因为它不是密钥，各环境相同）。
    只读一个文件的话会漏掉 issuer —— 而缺了它签出来的 token 验签直接失败，
    报错是 401，完全指不到"配置没读到"。

    这个文件是 YAML，但结构很浅，用行匹配比引 yaml 库更省事。找不到就明确报错，
    不要静默用默认值。
    """
    files = [DEV_YML, os.path.join(os.path.dirname(DEV_YML), "application.yml")]
    text = "\n".join(open(f, encoding="utf-8").read() for f in files if os.path.exists(f))
    out: dict[str, str] = {}
    for key, pattern in [
        ("jwt_secret", r"^\s*secret:\s*(\S+)\s*$"),
        ("hmac_key", r"^\s*hmac-key:\s*(\S+)\s*$"),
        ("issuer", r"^\s*issuer:\s*(\S+)\s*$"),
    ]:
        m = re.search(pattern, text, re.M)
        if not m:
            raise SystemExit(f"  ✗ 在 application*.yml 里找不到 {key}（{pattern}）")
        out[key] = m.group(1)
    return out


def b64u(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()


def mint_mp_token(secret_b64: str, issuer: str, uid: int, tenant_id: int,
                  roles: list[str], user_type: int) -> str:
    """手搓一个与 JwtService.issue 同构的 MP token（HS512）。"""
    key = base64.b64decode(secret_b64)
    now = int(time.time())
    header = {"alg": "HS512"}
    payload = {
        "iss": issuer,
        "sub": str(uid),
        "tenantId": tenant_id,
        "orgIds": [],
        "roles": roles,
        "userType": user_type,
        "iat": now,
        "exp": now + 3600,
    }
    signing = (b64u(json.dumps(header, separators=(",", ":")).encode()) + "."
               + b64u(json.dumps(payload, separators=(",", ":")).encode()))
    sig = hmac.new(key, signing.encode(), hashlib.sha512).digest()
    return signing + "." + b64u(sig)


def hmac_hex(key_b64: str, raw: str) -> str:
    """与 HmacIndex.hash 完全一致：trim + 转大写 → HmacSHA256 → hex。"""
    key = base64.b64decode(key_b64)
    return hmac.new(key, raw.strip().upper().encode("utf-8"), hashlib.sha256).hexdigest()


# ── HTTP ─────────────────────────────────────────────

def call(method: str, url: str, token: str | None = None,
         body: dict | None = None, timeout: int = 30) -> tuple[int, dict]:
    headers = {"Content-Type": "application/json; charset=utf-8"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    data = json.dumps(body, ensure_ascii=False).encode("utf-8") if body is not None else None
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        raw = e.read().decode("utf-8")
        try:
            return e.code, json.loads(raw)
        except json.JSONDecodeError:
            return e.code, {"raw": raw}


def rsa_encrypt(plain: str) -> str:
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import padding
    with open(PUB_KEY, "rb") as f:
        pub = serialization.load_pem_public_key(f.read())
    return base64.b64encode(pub.encrypt(plain.encode(), padding.PKCS1v15())).decode()


def login(base: str, username: str, password: str, tenant: str | None = None):
    payload = {"username": username, "password": rsa_encrypt(password)}
    if tenant:
        payload["tenantCode"] = tenant
    return call("POST", f"{base}/api/v1/auth/login", body=payload)


def match(base: str, token: str, name: str, id_card: str, task_id: int) -> dict:
    _, body = call("POST", f"{base}/api/v1/mp/roster/match", token,
                   {"realName": name, "idCardNo": id_card, "taskId": task_id})
    return body


def main() -> int:
    ap = argparse.ArgumentParser(description="门禁核验 E2E")
    ap.add_argument("--base", default=os.environ.get("XT_BASE", "http://localhost:8080"))
    ap.add_argument("--tenant", default=os.environ.get("XT_TENANT", "demo"))
    ap.add_argument("--password", default="Admin@123")
    args = ap.parse_args()
    base = args.base.rstrip("/")
    CTX["base"] = base

    print("=" * 72)
    print(f"门禁核验 E2E　base={base}")
    print("  （token 是用 dev 配置里的 JWT 密钥手搓的，见脚本头部说明）")
    print("=" * 72)

    suffix = uuid.uuid4().hex[:6]
    cfg = read_dev_config()

    # ── 0. 准备 ────────────────────────────────────────
    print("\n[0] 准备")
    status, body = login(base, "admin", args.password, args.tenant)
    if body.get("code") != 0:
        fail("管理员登录", f"code={body.get('code')} msg={body.get('msg')}")
        return 1
    admin_token = body["data"]["token"]
    CTX["token"] = admin_token
    tenant_id = (body["data"].get("tenant") or {}).get("id")
    admin_uid = (body["data"].get("user") or {}).get("id") or 2
    ok(f"管理员登录（租户 {tenant_id}，uid={admin_uid}）")

    # 手搓的学生 token
    student_uid = 990001
    stu_token = mint_mp_token(cfg["jwt_secret"], cfg["issuer"], student_uid, tenant_id,
                             ["STUDENT"], 4)
    other_uid = 990002
    other_token = mint_mp_token(cfg["jwt_secret"], cfg["issuer"], other_uid, tenant_id,
                              ["STUDENT"], 4)
    admin_mp = mint_mp_token(cfg["jwt_secret"], cfg["issuer"], admin_uid, tenant_id,
                            ["SCHOOL_ADMIN"], 2)

    _, probe = call("POST", f"{base}/api/v1/mp/roster/match", stu_token,
                    {"realName": "探针", "idCardNo": ID_A, "taskId": 999999})
    check("手搓的 token 能通过认证（不是 401）", probe.get("code") != 40101,
          f"code={probe.get('code')} msg={probe.get('msg')}")

    # 建任务
    status, body = call("POST", f"{base}/api/v1/tasks", admin_token, {
        "taskName": f"E2E门禁 {suffix}", "collectMode": 1, "batchNo": "2006",
        "majorName": "数控加工", "certType": 1,
        "startTime": "2026-09-28", "endTime": "2026-12-31",
    })
    if body.get("code") != 0:
        fail("建任务", f"code={body.get('code')} msg={body.get('msg')}")
        return 1
    task_id = body["data"]["id"]
    CREATED_TASK_IDS.append(task_id)
    ok(f"建任务（id={task_id}）")

    # 取一个班级
    status, body = call("GET", f"{base}/api/v1/orgs/tree", admin_token)
    classes: list[dict] = []

    def walk(nodes):
        for n in nodes or []:
            if n.get("orgType") == 3:
                classes.append(n)
            walk(n.get("children"))

    walk(body.get("data"))
    org_id = classes[0]["id"] if classes else None
    org_name = classes[0]["orgName"] if classes else None

    # 造名册（直接用 SQL：名册导入属 M2，这里只需要两行能匹配的数据）
    for name, idc, occ in [(NAME_A, ID_A, "车工"), (NAME_B, ID_B, "铣工")]:
        okk, out = sql(
            f"INSERT INTO collect_roster (tenant_id, task_id, org_id, real_name, "
            f"id_card_no, id_card_hash, id_card_mask, occupation, worktype, job_level, "
            f"education, is_student, phone) VALUES ("
            f"{tenant_id}, {task_id}, {org_id or 'NULL'}, '{name}', "
            f"'encrypted-placeholder', '{hmac_hex(cfg['hmac_key'], idc)}', "
            f"'{idc[:6]}******{idc[-4:]}', '{occ}', '数控{occ}', '四级', "
            f"'中等专科', 1, '138****8888')")
        if not okk:
            fail("造名册", out)
            return 1
    ok(f"造了 2 行名册（{NAME_A} / {NAME_B}）")

    clear_rate()

    # ── 1. 正常核验 ─────────────────────────────────────
    print("\n[1] 姓名 + 证件号完全匹配")
    r = match(base, stu_token, NAME_A, ID_A, task_id)
    check("核验通过", r.get("code") == 0, f"code={r.get('code')} msg={r.get('msg')}")
    d = r.get("data") or {}
    check("matched=true", d.get("matched") is True, f"{d.get('matched')}")
    check("返回 rosterId", isinstance(d.get("rosterId"), int), f"{d.get('rosterId')}")
    check("返回姓名", d.get("realName") == NAME_A, f"{d.get('realName')}")
    check("**证件号只回脱敏值**", d.get("idCardMask") and "******" in d["idCardMask"],
          f"idCardMask={d.get('idCardMask')}")
    check("返回里**没有任何字段含完整证件号**（防泄露）",
          ID_A not in json.dumps(d, ensure_ascii=False),
          f"响应里出现了明文证件号！{json.dumps(d, ensure_ascii=False)[:200]}")
    check("返回班级名", d.get("orgName") == org_name, f"{d.get('orgName')} 期望={org_name}")
    pf = d.get("prefill") or {}
    check("预填了报考职业（老师已填，学生只核对）", pf.get("occupation") == "车工",
          f"prefill={pf}")
    check("预填了报考工种与等级（锁定项）",
          pf.get("worktype") == "数控车工" and pf.get("jobLevel") == "四级", f"{pf}")
    check("needConfirmName 为 false（形近字确认属 FR-504c，本轮未实现）",
          d.get("needConfirmName") is False, f"{d.get('needConfirmName')}")

    # ── 2. 幂等 ─────────────────────────────────────────
    print("\n[2] 幂等：已绑自己的人再核验")
    r2 = match(base, stu_token, NAME_A, ID_A, task_id)
    check("重复核验仍成功（学生刷新页面不该看到「已被登记」）", r2.get("code") == 0,
          f"code={r2.get('code')} msg={r2.get('msg')}")

    # ── 3. 统一失败提示（防枚举的核心）────────────────────
    print("\n[3] 失败提示必须统一 —— 这是防身份证号枚举的底线")
    r_wrong_name = match(base, stu_token, "张三", ID_A, task_id)
    r_wrong_id = match(base, stu_token, NAME_A, "11010519491231002X", task_id)
    r_no_such = match(base, stu_token, NAME_A, "361029200912100015", task_id + 777)
    check("姓名不对 → 40007", r_wrong_name.get("code") == 40007,
          f"code={r_wrong_name.get('code')}")
    check("号码不在名册 → 40007", r_wrong_id.get("code") == 40007,
          f"code={r_wrong_id.get('code')}")
    check("**两者的提示文字完全一致**（分开就会变成枚举器：先用假姓名+目标号码试，"
          "若提示「姓名不对」就说明号码在名单里）",
          r_wrong_name.get("msg") == r_wrong_id.get("msg"),
          f"姓名错：{r_wrong_name.get('msg')} / 号码错：{r_wrong_id.get('msg')}")
    check("不存在的任务给的是另一类提示（那不属于身份信息不匹配）",
          r_no_such.get("code") == 40001, f"code={r_no_such.get('code')}")

    # ── 4. 已被他人绑定 ─────────────────────────────────
    print("\n[4] 同一名册项被别人核验")
    clear_rate()
    r_other = match(base, other_token, NAME_A, ID_A, task_id)
    check("另一账号核验同一行 → 40072 已被登记", r_other.get("code") == 40072,
          f"code={r_other.get('code')} msg={r_other.get('msg')}")

    # ── 5. 并发 ─────────────────────────────────────────
    print("\n[5] 并发：两个账号同时对同一行核验")
    clear_rate()
    # 用 B 那一行，让两个不同的 uid 同时冲
    import threading
    results: list[int] = []
    lock = threading.Lock()
    barrier = threading.Barrier(2)

    def hit(tok):
        barrier.wait()
        _, b = call("POST", f"{base}/api/v1/mp/roster/match", tok,
                    {"realName": NAME_B, "idCardNo": ID_B, "taskId": task_id})
        with lock:
            results.append(b.get("code"))

    t1 = threading.Thread(target=hit, args=(stu_token,))
    t2 = threading.Thread(target=hit, args=(other_token,))
    t1.start()
    t2.start()
    t1.join()
    t2.join()
    success_n = sum(1 for c in results if c == 0)
    bound_n = sum(1 for c in results if c == 40072)
    check(f"恰好一个成功、一个 40072（实际 {results}）", success_n == 1 and bound_n == 1,
          f"两个都成功会静默覆盖先绑的人，而先绑的学生之后再也进不去。结果={results}")
    okk, out = sql(f"SELECT count(*) FROM collect_roster WHERE task_id = {task_id} "
                   f"AND id_card_hash = '{hmac_hex(cfg['hmac_key'], ID_B)}' "
                   f"AND bind_user_id IS NOT NULL")
    check("库里那一行确实只被绑了一次", okk and out.strip() == "1", f"{out}")

    # ── 6. 格式非法的号码 ───────────────────────────────
    print("\n[6] 格式非法的号码")
    clear_rate()
    r_malformed = match(base, stu_token, NAME_A, "abcd", task_id)
    check("乱码号码 → 40007", r_malformed.get("code") == 40007, f"code={r_malformed.get('code')}")
    check("提示与「号码不在名册」**同样一致**（校验位的快速路径不泄露信息）",
          r_malformed.get("msg") == r_wrong_id.get("msg"),
          f"{r_malformed.get('msg')} / {r_wrong_id.get('msg')}")

    # ── 7. 限流 ─────────────────────────────────────────
    print("\n[7] 限流：同账号 10 次/小时")
    clear_rate()
    codes = []
    for _ in range(11):
        _, b = call("POST", f"{base}/api/v1/mp/roster/match", stu_token,
                    {"realName": "张三", "idCardNo": ID_A, "taskId": task_id})
        codes.append(b.get("code"))
    check("前 10 次是业务结果（40007）", codes[:10].count(40007) == 10, f"{codes}")
    check("第 11 次被限流（40071）", codes[10] == 40071, f"{codes}")
    clear_rate()
    _, after = call("POST", f"{base}/api/v1/mp/roster/match", stu_token,
                    {"realName": NAME_A, "idCardNo": ID_A, "taskId": task_id})
    check("清掉计数器后立刻恢复", after.get("code") == 0, f"code={after.get('code')}")

    # ── 8. 管理员也能核验（代学生核验是一期就要支持的能力）──
    print("\n[8] 管理员（代核验）")
    r_admin = match(base, admin_mp, "张三", ID_A, task_id)
    check("管理员也能调门禁（allowProxyUpload 默认开，代学生核验是支持的）",
          r_admin.get("code") in (0, 40007, 40072),
          f"code={r_admin.get('code')} msg={r_admin.get('msg')}")

    # ── 9. 清理 ─────────────────────────────────────────
    print("\n[9] 清理测试数据")
    cleanup_all()
    ok(f"已清理 {len(CREATED_TASK_IDS)} 个任务、{len(CREATED_USER_IDS)} 个账号")

    print("\n" + "=" * 72)
    print(f"通过 {len(PASSED)}　失败 {len(FAILED)}　共 {len(PASSED) + len(FAILED)}")
    if FAILED:
        print("\n失败项：")
        for name, detail in FAILED:
            print(f"  · {name}\n      {detail}")
    print("=" * 72)
    return 1 if FAILED else 0


if __name__ == "__main__":
    try:
        code = main()
    except BaseException:
        print("\n  ⚠ 脚本异常退出，正在清理…")
        cleanup_all()
        raise
    raise SystemExit(code)