#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""头像上传链路端到端验证（批次 A 分档 + B1 EXIF 归一）。

为什么走头像而不是提问图：
  - 提问图那条路要经过内容审核/提问入库才能拿到 key，链路长且易受其它因素干扰；
  - 头像强制 1:1 且 POST /user/profile/update 直接返回页面，便于立即取回 CDN URL。

验证两件事：
  A   分档生效：上传 1600x1600 的大图，头像档 maxSide=512 -> 存盘应为 512x512
      （改前是 1920，实测体积差 ~7 倍）
  B1  EXIF Orientation=6 归一：上传「像素 512x256 + EXIF 顺时针 90」的方形图，
      归一后仍应是 512x512（宽高互换但比例不变，故改用体积+尺寸双断言，
      并额外验证竖长条图 256x512 + Orientation=6 -> 应变成 512x256）

用法：
    python dev/devcheck-avatar.py [--base https://box-dev-api.337845818.best]
依赖：Pillow（仅用于本地生成测试图片）
清理：TamaBox ban --email <输出中的测试邮箱> --yes
"""
import argparse
import http.cookiejar
import io
import json
import re
import struct
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid

UA = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
      "(KHTML, like Gecko) Chrome/124.0 Safari/537.36")
RESULTS = []


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):  # noqa: D102
        return None


class Resp:
    def __init__(self, status, headers, body):
        self.status = status
        self.h = {k.lower(): v for k, v in headers.items()}
        self.body = body

    @property
    def loc(self):
        return self.h.get("location", "")

    @property
    def text(self):
        return self.body.decode("utf-8", "replace")


class Client:
    def __init__(self, base):
        self.base = base
        self.cj = http.cookiejar.CookieJar()
        self.op = urllib.request.build_opener(
            urllib.request.HTTPCookieProcessor(self.cj), NoRedirect())

    def _send(self, req, timeout=120):
        try:
            with self.op.open(req, timeout=timeout) as r:
                return Resp(r.status, dict(r.headers), r.read())
        except urllib.error.HTTPError as e:
            try:
                body = e.read()
            except Exception:  # noqa: BLE001
                body = b""
            return Resp(e.code, dict(e.headers), body)
        except Exception as e:  # noqa: BLE001
            return Resp(-1, {}, ("TRANSPORT_ERROR: %s" % e).encode())

    def req(self, method, path, data=None, content_type=None, timeout=120):
        url = path if path.startswith("http") else self.base + path
        headers = {"User-Agent": UA, "Referer": self.base + path.split("?")[0]}
        body = None
        if data is not None:
            body = data if isinstance(data, bytes) else urllib.parse.urlencode(data).encode()
            headers["Content-Type"] = content_type or "application/x-www-form-urlencoded"
        return self._send(urllib.request.Request(url, data=body, headers=headers, method=method),
                          timeout)

    def get(self, p, **kw):
        r = self.req("GET", p, **kw)
        return r if r.status != -1 else self.req("GET", p, **kw)

    def post(self, p, data=None, **kw):
        return self.req("POST", p, data=data, **kw)

    def get_binary(self, url):
        return self._send(urllib.request.Request(
            url, headers={"User-Agent": UA}, method="GET"), timeout=60)


def step(desc, ok, extra=""):
    RESULTS.append((desc, ok, extra))
    print("  [%s] %s%s" % ("PASS" if ok else "FAIL", desc, ("  (%s)" % extra) if extra else ""))
    return ok


def csrf_of(html):
    m = re.search(r'name="_csrf"\s+value="([^"]+)"', html)
    return m.group(1) if m else None


def captcha(c):
    j = json.loads(c.get("/api/captcha/generate").text)
    pid, q = j["data"]["id"], j["data"]["question"]
    m = re.match(r"\s*(\d+)\s*([+\-×])\s*(\d+)\s*=", q)
    a, op, b = int(m.group(1)), m.group(2), int(m.group(3))
    return "%s:%d" % (pid, {"+": a + b, "-": a - b, "×": a * b}[op])


def multipart(fields, files):
    boundary = "----tbav%s" % uuid.uuid4().hex
    buf = io.BytesIO()
    for k, v in fields.items():
        buf.write(("--%s\r\n" % boundary).encode())
        buf.write(('Content-Disposition: form-data; name="%s"\r\n\r\n' % k).encode())
        buf.write(("%s\r\n" % v).encode())
    for name, fn, ct, data in files:
        buf.write(("--%s\r\n" % boundary).encode())
        buf.write(('Content-Disposition: form-data; name="%s"; filename="%s"\r\n' % (name, fn)).encode())
        buf.write(("Content-Type: %s\r\n\r\n" % ct).encode())
        buf.write(data)
        buf.write(b"\r\n")
    buf.write(("--%s--\r\n" % boundary).encode())
    return buf.getvalue(), "multipart/form-data; boundary=%s" % boundary


def sniff(d):
    if len(d) >= 12 and d[:4] == b"RIFF" and d[8:12] == b"WEBP":
        return "webp", webp_dims(d)
    if d[:3] == b"\xff\xd8\xff":
        return "jpeg", jpeg_dims(d)
    if d[:8] == b"\x89PNG\r\n\x1a\n":
        return "png", (struct.unpack(">II", d[16:24]) if len(d) >= 24 else None)
    return "unknown", None


def webp_dims(d):
    c = d[12:16]
    if c == b"VP8 " and len(d) >= 30:
        return (struct.unpack("<H", d[26:28])[0] & 0x3FFF,
                struct.unpack("<H", d[28:30])[0] & 0x3FFF)
    if c == b"VP8L" and len(d) >= 25:
        b = struct.unpack("<I", d[21:25])[0]
        return (b & 0x3FFF) + 1, ((b >> 14) & 0x3FFF) + 1
    if c == b"VP8X" and len(d) >= 30:
        return ((d[24] | d[25] << 8 | d[26] << 16) + 1,
                (d[27] | d[28] << 8 | d[29] << 16) + 1)
    return None


def jpeg_dims(d):
    i = 2
    while i + 9 < len(d):
        if d[i] != 0xFF:
            i += 1
            continue
        m = d[i + 1]
        if m in (0xC0, 0xC1, 0xC2, 0xC3):
            h, w = struct.unpack(">HH", d[i + 5:i + 9])
            return (w, h)
        if m in (0xD8, 0xD9) or 0xD0 <= m <= 0xD7:
            i += 2
            continue
        i += 2 + struct.unpack(">H", d[i + 2:i + 4])[0]
    return None


def make_avatar(w, h, orientation=None):
    """生成一张带明显方位特征的方形头像。"""
    from PIL import Image, ImageDraw
    im = Image.new("RGB", (w, h), (30, 120, 200))
    d = ImageDraw.Draw(im)
    # 顶部一条粗红条：旋转 90 度后它会跑到侧面，方向一眼可辨
    d.rectangle([0, 0, w, max(2, h // 12)], fill=(230, 30, 30))
    for i in range(0, w, max(8, w // 20)):
        d.line([(i, h // 4), (i, h)], fill=(255, 255, 255), width=2)
    buf = io.BytesIO()
    if orientation:
        exif = im.getexif()
        exif[0x0112] = orientation
        im.save(buf, "JPEG", quality=90, exif=exif)
    else:
        im.save(buf, "JPEG", quality=90)
    return buf.getvalue()


def scrape_profile_fields(html):
    """从 profile 页抓取所有 input 的 name/value，保证回填必填项。

    UpdateProfile 的必填字段有十来个（昵称/介绍/5 个二维码颜色/点大小/透明度…），
    漏一个就会因 `valid:"required"` 被 Bind 挡下、头像根本不会走到上传分支。
    """
    fields = {}
    for m in re.finditer(r'<input\b[^>]*>', html, re.I):
        tag = m.group(0)
        nm = re.search(r'name="([^"]+)"', tag, re.I)
        if not nm:
            continue
        name = nm.group(1)
        if name in ("avatar", "background", "backgroundimage", "_csrf", "submit"):
            continue
        typ = re.search(r'type="([^"]+)"', tag, re.I)
        if typ and typ.group(1).lower() in ("checkbox", "radio"):
            # 只保留勾选态的，避免把「本想取消」的项又勾上
            if not re.search(r'\bchecked\b', tag, re.I):
                continue
        val = re.search(r'value="([^"]*)"', tag, re.I)
        fields[name] = val.group(1) if val else ""
    # select 走第一个非空 option
    for m in re.finditer(r'<select\b[^>]*name="([^"]+)"[^>]*>(.*?)</select>', html, re.I | re.S):
        name, body = m.group(1), m.group(2)
        opts = re.findall(r'<option\b([^>]*)>', body, re.I)
        chosen = ""
        for o in opts:
            vm = re.search(r'value="([^"]*)"', o)
            v = vm.group(1) if vm else ""
            if re.search(r'\bselected\b', o, re.I):
                chosen = v
                break
            if not chosen:
                chosen = v
        fields[name] = chosen
    # textarea
    for m in re.finditer(r'<textarea\b[^>]*name="([^"]+)"[^>]*>(.*?)</textarea>', html, re.I | re.S):
        fields[m.group(1)] = re.sub(r"<[^>]+>", "", m.group(2)).strip()
    return fields


def upload_avatar(c, data, fname="avatar.jpg"):
    page = c.get("/user/profile")
    if page.status != 200:
        return None, "profile HTTP %s" % page.status
    fields = scrape_profile_fields(page.text)
    fields["_csrf"] = csrf_of(page.text)
    body, ct = multipart(fields, [("avatar", fname, "image/jpeg", data)])
    res = c.post("/user/profile/update", data=body, content_type=ct)
    # 成功时是 302 跳回 /user/profile，头像 URL 在跳转后的页面里
    if res.status == 302 and res.loc:
        res = c.get(res.loc if res.loc.startswith("/") else "/user/profile")
    if res.status != 200:
        return None, "update 后页面 HTTP %s" % res.status
    # 从页面里取头像 URL：<img src="https://r2.../picture/..."> 或 og:image
    urls = re.findall(r'https://r2\.[^"\'\s<>]+/picture/[0-9]+/\d{2}/\d{2}/[0-9a-f]+', res.text)
    if not urls:
        m = re.search(r'(https://r2\.[^"\'\s<>]+/picture/[^"\'\s<>]+)', res.text)
        urls = [m.group(1)] if m else []
    if urls:
        return urls[-1], None
    # 没匹配到：落盘返回页并打印可疑片段，便于定位是「上传被拒」还是「正则没覆盖」
    dump = "C:/Users/Tamakyi Gekka/AppData/Local/Temp/tbav_profile.html"
    with open(dump, "wb") as f:
        f.write(res.body)
    alerts = re.findall(r'class="uk-alert[^"]*"[^>]*>(.*?)</div>', res.text, re.S)
    hint = " | ".join(re.sub(r"\s+", " ", re.sub(r"<[^>]+>", " ", a)).strip()[:80]
                      for a in alerts[:3])
    return None, "未匹配到 picture URL（已存 %s）%s%s" % (dump, (" 提示:" + hint) if hint else "",
                                                        (" HTTP %s" % res.status))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="https://box-dev-api.337845818.best")
    args = ap.parse_args()
    base = args.base.rstrip("/")
    pwd = "TbAv!2026"
    ts = time.strftime("%m%d%H%M%S")
    email, domain = "tbav%s@example.com" % ts, "tbav%s" % ts
    print("== 账号：%s / %s ==" % (email, domain))

    # 方形 1600x1600：验证 A 分档（应存成 512x512）
    square = make_avatar(1600, 1600)
    # EXIF Orientation=6（顺时针 90）：像素存成横的 1536x512，归一后显示为竖的 512x1536，
    # 再经头像档限边 512 -> 512x512。
    # 注意必须用「像素横躺」的样本：头像有 1:1 比例校验，若像素本身就是竖的会被先拒掉，
    # 那样就测不到 EXIF 分支了。
    exif_shot = make_avatar(1536, 512, orientation=6)
    print("   本地样本 square   1600x1600%7d bytes" % len(square))
    print("   本地样本 exifshot 1536x512 (EXIF=6)%d bytes" % len(exif_shot))

    c = Client(base)
    r = c.get("/register")
    reg = c.post("/register", {
        "email": email, "domain": domain, "name": "TBAv%s" % ts,
        "password": pwd, "repeat_password": pwd,
        "g-recaptcha-response": captcha(c), "language": "zh-CN",
        "_csrf": csrf_of(r.text),
    })
    if not step("注册", reg.status == 302, "HTTP %s -> %s" % (reg.status, reg.loc)):
        return 1

    c = Client(base)
    r = c.get("/login")
    lg = c.post("/login", {"email": email, "password": pwd,
                           "g-recaptcha-response": captcha(c), "_csrf": csrf_of(r.text)})
    if not step("登录", lg.status == 302, "HTTP %s -> %s" % (lg.status, lg.loc)):
        return 1

    # ---- A 分档 ----
    url, err = upload_avatar(c, square)
    if not step("上传方形大头像（1600x1600）", url is not None, err or (url or "")):
        return 1
    img = c.get_binary(url)
    st, data = img.status, img.body
    if not step("拉回头像", st == 200 and bool(data), "HTTP %s" % st):
        return 1
    kind, dims = sniff(data)
    print("       存储结果：%s %s %d bytes" % (kind, dims, len(data)))
    step("A 分档生效：大图被缩到头像档 512（改前是 1920）",
         dims == (512, 512), "实际 %s" % (dims,))

    # ---- B1 EXIF ----
    # 关键：两次上传的 URL 必须不同，否则说明第二次没生效（回退到了旧头像）。
    url2, err2 = upload_avatar(c, exif_shot)
    if not step("上传 EXIF 横躺头像（1536x512, Orientation=6）", bool(url2), err2 or (url2 or "")):
        return 0
    if url2 == url:
        step("B1 EXIF 归一（第二次上传未生效，URL 未变）", False,
             "两次上传返回同一 URL，无法区分；请检查是否被 1:1 校验拒绝")
        return 0
    img2 = c.get_binary(url2)
    kind2, dims2 = sniff(img2.body)
    print("       存储结果：%s %s %d bytes" % (kind2, dims2, len(img2.body)))
    # 归一后：像素 1536x512 -> 旋转成 512x1536 -> 再限边 512 -> 512x512
    step("B1 EXIF Orientation=6 归一（1536x512 -> 512x1536 -> 512x512）",
         dims2 == (512, 512), "实际 %s" % (dims2,))
    # 体积对照：未归一时 1536x512 的 WebP 明显更大
    step("B1 归一后体积已按 512 档压过（应与方形样本同量级）",
         len(img2.body) < 6000, "%d bytes" % len(img2.body))

    print("\n== 结果：%d/%d 通过 ==" % (sum(1 for _, ok, _ in RESULTS if ok), len(RESULTS)))
    print("清理：TamaBox ban --email %s --yes" % email)
    return 0 if all(ok for _, ok, _ in RESULTS) else 1


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