#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
OCS 5.0「图片题优化」对抗性 E2E 测试（会真实调用 DeepSeek，约 6 次，费用很小）

覆盖 Lead 要求的 6 组对抗场景：
  a) markers 数量 > images 数量时的降级（[图片3] 无处可放）
  b) images 里有非字符串 / 空值 / 非法项
  c) suggestion_title/options 缺失但 images 有值
  d) 大 payload（2 张约 1MB base64 图片 + 1 张题干图）能正常答题，且日志不爆炸
  e) 纯文字题：images=[]、suggestion_* 为空字符串
  f) 旧版 payload（无 images 字段）：仍能走 URL 下载视觉链路答题
  g) 额外（真实 OCS 细节）：GM 回退路径会把图片原始字节的 mime 前缀统一改写成
     data:image/png;base64,...（例如 JPEG 字节+PNG 前缀），验证 DeepSeek 仍能识别。

除真实接口断言外，每组都在本地先用 ai_client 的纯函数做「消息构建预检」（不花钱），
并全程捕获 "ocs_ai" 日志，检查 base64 不会进入日志（日志脱敏）。

运行（项目根目录）：
    $env:PYTHONIOENCODING='utf-8'
    py tests/test_ocs5_adversarial_e2e.py                # 离线预检 + 真实 E2E（约 6 次 API）
    py tests/test_ocs5_adversarial_e2e.py --offline-only # 只做离线预检，不联网不花钱
"""

import base64
import io
import logging
import os
import random
import re
import sys
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
os.chdir(ROOT)

OFFLINE_ONLY = "--offline-only" in sys.argv

FAILS = []
RISKS = []


def record(name: str, ok: bool, detail: str = "") -> bool:
    print(("  [PASS] " if ok else "  [FAIL] ") + name + (f" | {detail}" if detail else ""))
    if not ok:
        FAILS.append(name)
    return ok


def risk(name: str, detail: str = "") -> None:
    print(f"  [RISK] {name}" + (f" | {detail}" if detail else ""))
    RISKS.append(f"{name} | {detail}")


def info(msg: str) -> None:
    print(f"  [INFO] {msg}")


# ────────────────────────── 图片生成 ──────────────────────────
FONT_CANDIDATES = [
    r"C:\Windows\Fonts\arial.ttf",
    r"C:\Windows\Fonts\arialbd.ttf",
    r"C:\Windows\Fonts\segoeui.ttf",
]


def _load_font(size: int):
    from PIL import ImageFont

    for path in FONT_CANDIDATES:
        if os.path.isfile(path):
            try:
                return ImageFont.truetype(path, size)
            except Exception:
                continue
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


def make_png(text: str, width: int = 700, height: int = 130, font_size: int = 46) -> bytes:
    from PIL import Image, ImageDraw

    img = Image.new("RGB", (width, height), "white")
    draw = ImageDraw.Draw(img)
    draw.text((30, height // 2 - font_size // 2 - 6), text, fill="black", font=_load_font(font_size))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def make_jpg(text: str, width: int = 700, height: int = 130, font_size: int = 44) -> bytes:
    from PIL import Image, ImageDraw

    img = Image.new("RGB", (width, height), "white")
    draw = ImageDraw.Draw(img)
    draw.text((30, height // 2 - font_size // 2 - 6), text, fill="black", font=_load_font(font_size))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=90)
    return buf.getvalue()


def to_data_url(png: bytes) -> str:
    return "data:image/png;base64," + base64.b64encode(png).decode("ascii")


def make_noise_png(size: int = 512, seed: int = 20261102) -> bytes:
    """生成随机噪声 PNG（用来撑大 payload；内容是噪声，不期望模型读出含义）"""
    from PIL import Image

    rnd = random.Random(seed)
    raw = bytes(rnd.getrandbits(8) for _ in range(size * size * 3))
    img = Image.frombytes("RGB", (size, size), raw)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


# ────────────────────────── 答案归一化 ──────────────────────────
def normalize_answer(raw) -> str:
    if raw is None:
        return ""
    text = str(raw).strip().upper()
    if not text:
        return ""
    letters = re.findall(r"(?<![A-Z])([A-H])(?![A-Z])", text)
    if letters:
        uniq = list(dict.fromkeys(letters))
        return "#".join(uniq) if len(uniq) > 1 else uniq[0]
    return re.sub(r"[^A-Z]", "", text)


# ────────────────────────── 日志捕获 ──────────────────────────
class LogCapture(logging.Handler):
    def __init__(self):
        super().__init__(level=logging.DEBUG)
        self.messages = []

    def emit(self, record):
        try:
            self.messages.append(record.getMessage())
        except Exception:
            self.messages.append("<log error>")

    def reset(self):
        self.messages = []


CAPTURE = LogCapture()


def check_logs(stage: str, *, expect_images: int = None) -> bool:
    """检查日志脱敏：不含 base64 内容、没有爆炸长行。"""
    messages = list(CAPTURE.messages)
    joined = "\n".join(messages)
    max_line = max((len(m) for m in messages), default=0)
    ok = True
    if "iVBORw0KGgo" in joined:
        record(f"{stage}：日志不包含图片 base64", False, "发现 PNG base64 签名")
        ok = False
    if "data:image/png;base64," in joined:
        record(f"{stage}：日志不包含 data URL", False, "发现 data:image/png;base64,")
        ok = False
    if max_line >= 2000:
        record(f"{stage}：日志单行长度 < 2000", False, f"max={max_line}")
        ok = False
    body_logs = [m for m in messages if "POST 请求体" in m]
    if expect_images is not None:
        want = f"<{expect_images} 张 base64 图片>"
        if not any(want in m for m in body_logs):
            record(f"{stage}：日志中 images 显示为占位文案", False, f"未找到 {want!r}")
            ok = False
    if ok:
        record(f"{stage}：日志脱敏正常（max_line={max_line}）", True)
    return ok


# ────────────────────────── 离线预检 ──────────────────────────
def offline_checks() -> dict:
    """不花钱的消息构建预检，返回后续 E2E 需要的图片素材。"""
    print("\n[离线预检] ai_client 消息构建 / 容错（不调 DeepSeek）")
    from ai_client import (
        QuestionInfo,
        _iter_marker_segments,
        _ocs_vision_messages,
        _prepare_ocs_images,
    )

    q_png = make_png("Solve for x: 3x + 5 = 20")
    q_data_url = to_data_url(q_png)
    digit_pngs = {d: make_png(d, 220, 150, 64) for d in ("3", "5", "7", "9")}

    # (a) markers 数量 > images 数量
    qa = QuestionInfo(
        title="原始标题（含坏 URL） https://p.ananas.chaoxing.com/not-found.png",
        type="single",
        options="A. 3\nB. 5\nC. 7\nD. 9",
        images=[q_data_url],
        suggestion_title="请解方程： [图片1] 另外 [图片3] ",
        suggestion_options="",
    )
    prepared_a = _prepare_ocs_images(qa.images)
    msgs_a = _ocs_vision_messages(qa, prepared_a)
    blocks_a = msgs_a[1]["content"]
    img_blocks_a = [b for b in blocks_a if b.get("type") == "image_url"]
    text_a = " ".join(b.get("text", "") for b in blocks_a if b.get("type") == "text")
    record("a) markers>images：有效图仍被正确放入消息", len(img_blocks_a) == 1,
           f"image_blocks={len(img_blocks_a)}")
    record("a) markers>images：越界标记 [图片3] 退化为文本（不崩溃）",
           "[图片3]" in text_a and "[图片1]" not in text_a, text_a[:80])
    record("a) markers>images：越界标记切分结果", 
           [(k, v) for k, v in _iter_marker_segments("x [图片3] y") if k == "image"] == [("image", 3)])

    # (b) images 含非字符串 / 空值
    raw_images_b = [None, "", "undefined", 123, {"oops": 1}, q_data_url]
    filtered_b = _prepare_ocs_images(raw_images_b)
    record("b) images 含非字符串/空值：过滤后只留有效图", filtered_b == [q_data_url],
           f"len={len(filtered_b)}")
    qb = QuestionInfo(
        title="原始标题",
        type="single",
        options="A. 3\nB. 5\nC. 7\nD. 9",
        images=raw_images_b,
        suggestion_title="请解方程： [图片1] ",
        suggestion_options="",
    )
    msgs_b = _ocs_vision_messages(qb, filtered_b)
    record("b) 过滤后 [图片1] 仍能映射到有效图",
           [b for b in msgs_b[1]["content"] if b.get("type") == "image_url"] != [])

    # (c) suggestion_* 缺失但 images 有值
    # R1 加固后：图片不再静默丢弃，会以「题目图片：」兜底附到消息里
    qc = QuestionInfo(
        title="Solve for x: 3x + 5 = 20",
        type="single",
        options="A. 3\nB. 5\nC. 7\nD. 9",
        images=[q_data_url],
        suggestion_title="",
        suggestion_options="",
    )
    msgs_c = _ocs_vision_messages(qc, _prepare_ocs_images(qc.images))
    img_blocks_c = [b for b in msgs_c[1]["content"] if b.get("type") == "image_url"]
    text_c = " ".join(b.get("text", "") for b in msgs_c[1]["content"] if b.get("type") == "text")
    record(
        "c) suggestion_* 缺失：不崩溃，图片走兜底附到消息中",
        len(img_blocks_c) == 1 and "题目图片" in text_c,
        f"image_blocks={len(img_blocks_c)}",
    )

    # c-风险变体：title 是裸 URL、suggestion_* 缺失、只收到 images
    # 这种 payload 只会出现在「题库配置只引用了 ${images} 而没引用 suggestion_*」的错误配置下。
    qc_risk = QuestionInfo(
        title="https://p.ananas.chaoxing.com/star3/origin/q.png",
        type="single",
        options="A. 3\nB. 5\nC. 7\nD. 9",
        images=[q_data_url],
        suggestion_title="",
        suggestion_options="",
    )
    msgs_risk = _ocs_vision_messages(qc_risk, _prepare_ocs_images(qc_risk.images))
    blocks_risk = msgs_risk[1]["content"]
    risk_imgs = [b for b in blocks_risk if b.get("type") == "image_url"]
    risk_text = " ".join(b.get("text", "") for b in blocks_risk if b.get("type") == "text")
    if not risk_imgs:
        risk(
            "suggestion_* 缺失且 title 只有裸 URL 时，images 被静默忽略 → 题目信息丢失（不崩溃）",
            f"repro: images=[data_url], suggestion_title='', title='<图片URL>'；消息中 image_blocks=0，text={risk_text[:60]!r}",
        )
    else:
        record("c-风险变体：suggestion_* 缺失时会回退使用 images", True)

    # (e) 纯文字题
    qe = QuestionInfo(title="2 + 3 = ?", type="single", options="A. 4\nB. 5\nC. 6\nD. 7",
                      images=[], suggestion_title="", suggestion_options="")
    msgs_e = _ocs_vision_messages(qe, [])
    record("e) 纯文字题：消息中无图片块",
           [b for b in msgs_e[1]["content"] if b.get("type") == "image_url"] == [])

    # (f) 旧 payload 的 URL 抽取仍正常
    from ai_client import _extract_image_urls

    record("f) 旧 payload：能从 title 抽出本地图片 URL",
           _extract_image_urls("题目 https://127.0.0.1:1/x.png") == ["https://127.0.0.1:1/x.png"])

    # (g) 真实 OCS 细节：GM_xmlhttpRequest 回退路径下，normalizeToPngDataUrl 只改前缀不转码，
    # 可能出现「JPEG 字节 + data:image/png 前缀」。中间件原样透传，需确认可被视觉模型接受。
    from ai_client import _normalize_ocs_image

    jpeg_bytes = make_jpg("Solve for x: 3x + 5 = 20")
    mismatched_jpeg_url = "data:image/png;base64," + base64.b64encode(jpeg_bytes).decode("ascii")
    record("g) 中间件接受「JPEG 字节 + PNG 前缀」的 data URL（OCS GM 回退形态）",
           _normalize_ocs_image(mismatched_jpeg_url) == mismatched_jpeg_url)

    # (d) 大图片素材：2 张约 1MB(base64) 噪声图 + 题干图
    noise1 = make_noise_png(512, seed=111)
    noise2 = make_noise_png(512, seed=222)
    noise1_url = to_data_url(noise1)
    noise2_url = to_data_url(noise2)
    info(f"大图素材：noise1 {len(noise1) // 1024}KB(png) / base64 {len(noise1_url) // 1024}KB，"
         f"noise2 {len(noise2) // 1024}KB(png) / base64 {len(noise2_url) // 1024}KB")
    record("d) 大图 base64 每张约 ≥ 700KB（构造大 payload 前提）",
           len(noise1_url) >= 700_000 and len(noise2_url) >= 700_000,
           f"{len(noise1_url)} / {len(noise2_url)}")

    return {
        "q_png": q_png,
        "q_data_url": q_data_url,
        "digit_pngs": digit_pngs,
        "noise1_url": noise1_url,
        "noise2_url": noise2_url,
        "mismatched_jpeg_url": mismatched_jpeg_url,
    }


# ────────────────────────── 本地图片服务（f 场景用） ──────────────────────────
class _ImageHandler(BaseHTTPRequestHandler):
    png = b""
    requests = []

    def log_message(self, fmt, *args):
        return

    def do_GET(self):  # noqa: N802
        referer = self.headers.get("Referer") or ""
        if self.path.split("?")[0] == "/q.png":
            self.requests.append({"path": "/q.png", "referer": referer})
            self.send_response(200)
            self.send_header("Content-Type", "image/png")
            self.send_header("Content-Length", str(len(self.png)))
            self.end_headers()
            self.wfile.write(self.png)
        else:
            self.requests.append({"path": self.path, "referer": referer})
            self.send_response(404)
            self.end_headers()


def start_image_server(png: bytes):
    handler_cls = type("ImageHandler", (_ImageHandler,), {"png": png, "requests": []})
    server = ThreadingHTTPServer(("127.0.0.1", 0), handler_cls)
    server.daemon_threads = True
    thread = threading.Thread(target=server.serve_forever, name="adv-image-server", daemon=True)
    thread.start()
    return server, handler_cls, server.server_address[1]


# ────────────────────────── E2E 对抗场景 ──────────────────────────
def e2e_checks(images: dict) -> None:
    print("\n[E2E] 真实调用 POST /api/search（会消耗少量 DeepSeek 额度）")

    from fastapi.testclient import TestClient

    from server import app

    q_data_url = images["q_data_url"]
    digit = {k: to_data_url(v) for k, v in images["digit_pngs"].items()}

    scenarios = []

    # a) markers 数量 > images 数量
    scenarios.append(
        {
            "id": "a",
            "label": "markers([图片3]) 超出 images 数量",
            "payload": {
                "title": "原始标题（含坏 URL） https://p.ananas.chaoxing.com/star3/origin/not-found.png",
                "type": "single",
                "options": "A. 3\nB. 5\nC. 7\nD. 9",
                "images": [q_data_url],
                "suggestion_title": "请解方程： [图片1] 另外 [图片3] ",
                "suggestion_options": "",
            },
            "expect_images_log": 1,
        }
    )

    # b) images 含非字符串 / 空值
    scenarios.append(
        {
            "id": "b",
            "label": "images 含 None/空串/数字/dict/undefined",
            "payload": {
                "title": "原始标题",
                "type": "single",
                "options": "A. 3\nB. 5\nC. 7\nD. 9",
                "images": [None, "", "undefined", 123, {"oops": 1}, q_data_url],
                "suggestion_title": "请解方程： [图片1] ",
                "suggestion_options": "",
            },
            "expect_images_log": 6,
        }
    )

    # c) suggestion_title/options 缺失但 images 有值（题面本身是纯文字，应能答对）
    scenarios.append(
        {
            "id": "c",
            "label": "suggestion_* 缺失但 images 有值（按纯文字题退化）",
            "payload": {
                "title": "Solve for x: 3x + 5 = 20",
                "type": "single",
                "options": "A. 3\nB. 5\nC. 7\nD. 9",
                "images": [q_data_url],
            },
            "expect_images_log": 1,
        }
    )

    # e) 纯文字题（OCS 正常形态：images=[]、suggestion_* 空串）
    scenarios.append(
        {
            "id": "e",
            "label": "纯文字题 images=[] suggestion_*=''",
            "payload": {
                "title": "2 + 3 = ?",
                "type": "single",
                "options": "A. 4\nB. 5\nC. 6\nD. 7",
                "images": [],
                "suggestion_title": "",
                "suggestion_options": "",
            },
            "expect_images_log": 0,
        }
    )

    # d) 大 payload：2 张 ~1MB base64 噪声图 + 1 张题干图（[图片3]）
    large_payload = {
        "title": "请阅读题目图片后作答",
        "type": "single",
        "options": "A. 3\nB. 5\nC. 7\nD. 9",
        "images": [images["noise1_url"], images["noise2_url"], q_data_url],
        "suggestion_title": "请阅读题目： [图片3] ",
        "suggestion_options": "",
    }
    import json as _json

    est_size = len(_json.dumps(large_payload))
    info(f"d) 大 payload 预计 JSON 大小 ≈ {est_size / 1048576:.2f} MB")
    record("d) 大 payload 体积 > 2MB", est_size > 2 * 1024 * 1024, f"{est_size} bytes")
    scenarios.append(
        {
            "id": "d",
            "label": "大 payload（2×~1MB base64 + 题干图）",
            "payload": large_payload,
            "expect_images_log": 3,
        }
    )

    # g) 真实 OCS 细节：JPEG 字节 + data:image/png 前缀（GM_xmlhttpRequest 回退路径的形态）
    scenarios.append(
        {
            "id": "g",
            "label": "JPEG 字节 + PNG 前缀的 data URL（OCS GM 回退路径）",
            "payload": {
                "title": "请阅读题目图片：",
                "type": "single",
                "options": "A. 3\nB. 5\nC. 7\nD. 9",
                "images": [images["mismatched_jpeg_url"]],
                "suggestion_title": "请阅读题目图片： [图片1] ",
                "suggestion_options": "",
            },
            "expect_images_log": 1,
        }
    )

    # f) 旧版 payload：无 images/suggestion 字段，title 里是图片 URL → URL 下载视觉链路
    server, handler_cls, port = start_image_server(images["q_png"])
    q_url = f"http://127.0.0.1:{port}/q.png"
    scenarios.append(
        {
            "id": "f",
            "label": "旧版 payload（无 images）→ URL 下载链路",
            "payload": {
                "title": f"{q_url}",
                "type": "single",
                "options": "A. 3\nB. 5\nC. 7\nD. 9",
            },
            # 旧 payload 没有 images 字段，日志里不会出现 "<N 张 base64 图片>"
            "expect_images_log": None,
        }
    )

    results = []
    ocs_logger = logging.getLogger("ocs_ai")
    ocs_logger.addHandler(CAPTURE)
    try:
        # 必须用 with 触发 FastAPI lifespan
        with TestClient(app) as client:
            for sc in scenarios:
                print(f"\n  [{sc['id']}] {sc['label']}")
                CAPTURE.reset()
                try:
                    resp = client.post("/api/search", json=sc["payload"])
                    data = resp.json() if resp.status_code == 200 else {}
                except Exception as exc:
                    import traceback

                    traceback.print_exc()
                    results.append({"id": sc["id"], "ok": False, "error": repr(exc)})
                    print(f"    POST 异常: {exc!r}")
                    continue

                answer = str(data.get("answer", "")) if isinstance(data, dict) else ""
                norm = normalize_answer(answer)
                ok_http = resp.status_code == 200 and data.get("code") == 1
                ok = ok_http and norm == "B"
                print(f"    HTTP {resp.status_code} code={data.get('code')!r} "
                      f"answer={answer!r} 归一化={norm!r} 期望='B' → {'PASS' if ok else 'FAIL'}")
                if not ok_http and data.get("msg"):
                    print(f"    msg: {str(data.get('msg'))[:300]}")

                log_ok = check_logs(f"[{sc['id']}]", expect_images=sc.get("expect_images_log"))
                results.append({"id": sc["id"], "ok": ok and log_ok, "answer": answer, "norm": norm,
                                "log_ok": log_ok})
    finally:
        ocs_logger.removeHandler(CAPTURE)
        server.shutdown()
        server.server_close()

    # f 场景：旧链路应带 Referer 请求本地图片
    reqs = list(handler_cls.requests)
    print("\n  [f] 本地图片服务请求记录：")
    for r in reqs:
        print(f"      GET {r['path']} | Referer={r['referer']!r}")
    record("f) 旧链路下载图片时带了 chaoxing Referer",
           any("chaoxing" in (r["referer"] or "").lower() for r in reqs),
           f"requests={len(reqs)}")

    passed = sum(1 for r in results if r["ok"])
    print(f"\n  E2E 对抗场景: PASS {passed}/{len(results)}")
    for r in results:
        if not r["ok"]:
            record(f"场景 {r['id']} 通过", False, f"answer={r.get('answer')!r} log_ok={r.get('log_ok')}")


# ────────────────────────── main ──────────────────────────
def main() -> int:
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        sys.stderr.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass

    print("=" * 74)
    print("OCS 5.0 图片题优化 · 对抗性测试" + ("（仅离线预检）" if OFFLINE_ONLY else "（离线预检 + 真实 E2E）"))
    print("=" * 74)

    images = offline_checks()
    if not OFFLINE_ONLY:
        e2e_checks(images)

    print("\n" + "=" * 74)
    if RISKS:
        print(f"风险提示 {len(RISKS)} 条（不阻断，见上方 [RISK]）：")
        for r in RISKS:
            print(f"  - {r}")
    if FAILS:
        print(f"总体: FAIL ❌ ({len(FAILS)} 项失败)")
        for name in FAILS:
            print(f"  - {name}")
        print("=" * 74)
        return 1
    print("总体: PASS ✅" + ("（离线预检）" if OFFLINE_ONLY else f"（对抗场景全部通过）"))
    print("=" * 74)
    return 0


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