# -*- coding: utf-8 -*-
"""
OCS 5.0「图片题优化」链路测试

覆盖：
  * 题库配置是否满足 OCS 5.0 图片支持检测（POST + ${images}/${suggestion_title}/${suggestion_options}）
  * base64 图片数组 + [图片N] 标记的离线消息构建
  * 真实 /api/search E2E（title 图片 / 选项图片 / 混合多图 / 纯文字回归）
  * 旧版 payload（无 images 字段）兼容

运行（项目根目录）：
  $env:PYTHONIOENCODING='utf-8'
  py tests/test_ocs_images_e2e.py              # 全部（真实调用 DeepSeek，少量费用）
  py tests/test_ocs_images_e2e.py --offline-only  # 只跑离线部分
"""

import base64
import io
import os
import re
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
os.chdir(ROOT)

from PIL import Image, ImageDraw, ImageFont  # noqa: E402

FONT_PATH = r"C:\Windows\Fonts\arial.ttf"


def font(size: int):
    try:
        return ImageFont.truetype(FONT_PATH, size)
    except Exception:
        return ImageFont.load_default()


def make_image(text: str, width: int = 760, height: int = 140, size: int = 46) -> str:
    """生成一张白底黑字 PNG 并返回 data URL"""
    img = Image.new("RGB", (width, height), "white")
    draw = ImageDraw.Draw(img)
    draw.text((30, height // 2 - size // 2 - 6), text, fill="black", font=font(size))
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/png;base64,{b64}"


FAILS = []


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


# ═══════════════ 离线部分 ═══════════════


def offline_checks() -> None:
    print("\n[离线] 图片题优化辅助函数 / 题库配置检测")

    from ai_client import (
        QuestionInfo,
        _iter_marker_segments,
        _ocs_vision_messages,
        _prepare_ocs_images,
    )

    segs = _iter_marker_segments("求值 [图片1] 与 [图片2] 的关系")
    record(
        "标记切分",
        [s[0] for s in segs] == ["text", "image", "text", "image", "text"]
        and segs[1][1] == 1
        and segs[3][1] == 2,
        str(segs),
    )

    png = make_image("1 + 1 = ?")
    imgs = _prepare_ocs_images([png, "", "undefined", None, 123, "data:image/jpeg;base64,AAAA"])
    record("图片筛选", len(imgs) == 2 and imgs[0] == png, f"len={len(imgs)}")
    record("超量截断", len(_prepare_ocs_images([png] * 30)) == 20)

    q = QuestionInfo(
        title="原始标题 https://p.ananas.chaoxing.com/x.png",
        type="single",
        options="A. [图片2]\nB. [图片3]",
        images=[png, png, png],
        suggestion_title="求 x 的值 [图片1]",
        suggestion_options="A. [图片2]\nB. [图片3]",
    )
    messages = _ocs_vision_messages(q, [png, png, png])
    blocks = messages[1]["content"]
    types = [b["type"] for b in blocks]
    texts = " ".join(b.get("text", "") for b in blocks if b["type"] == "text")
    record("多模态图片块数量", types.count("image_url") == 3, f"images={types.count('image_url')}")
    record("标题优先用 suggestion_title", "求 x 的值" in texts and "ananas" not in texts)
    record("选项字母保留", "A." in texts and "B." in texts, texts[:80])
    record("末条为作答要求", blocks[-1]["type"] == "text" and "只回答一个正确选项" in blocks[-1]["text"])

    # 兜底 1：有 images 但文本里没有任何 [图片N] 标记（配置只引用 ${images} 的异常情况），图片不能丢
    q_no_mark = QuestionInfo(
        title="题目图片：https://p.ananas.chaoxing.com/x.png",
        type="single",
        options="A. 1\nB. 2",
        images=[png, png, png],
        suggestion_title="",
        suggestion_options="",
    )
    blocks_no_mark = _ocs_vision_messages(q_no_mark, [png, png, png])[1]["content"]
    no_mark_text = " ".join(b.get("text", "") for b in blocks_no_mark if b["type"] == "text")
    record(
        "无标记时图片兜底",
        [b["type"] for b in blocks_no_mark].count("image_url") == 3
        and "题目图片" in no_mark_text,
        f"images={[b['type'] for b in blocks_no_mark].count('image_url')}",
    )

    # 兜底 2：标记只用到部分图片，剩余图片也要追加给模型
    q_partial = QuestionInfo(
        title="见 [图片1]",
        type="single",
        options="",
        images=[png, png, png],
        suggestion_title="见 [图片1]",
        suggestion_options="",
    )
    blocks_partial = _ocs_vision_messages(q_partial, [png, png, png])[1]["content"]
    record(
        "未使用图片追加",
        [b["type"] for b in blocks_partial].count("image_url") == 3
        and any("其余图片" in (b.get("text") or "") for b in blocks_partial if b["type"] == "text"),
    )

    # OCS 5.0 官方检测逻辑（isAnswererWrapperSupportImageOptimize 的等价实现）
    from generate_ocs_config import generate_config

    def ocs_supports_image(wrappers: list) -> bool:
        placeholders = ("${images}", "${suggestion_title}", "${suggestion_options}")
        env_ref = re.compile(r"(^|[^\w$])(images|suggestion_title|suggestion_options)(?![\w$])")
        for wrapper in wrappers:
            if str(wrapper.get("method", "get")).lower() != "post":
                continue
            for value in (wrapper.get("data") or {}).values():
                if isinstance(value, str) and any(p in value for p in placeholders):
                    return True
                if isinstance(value, dict) and env_ref.search(str(value.get("handler", ""))):
                    return True
        return False

    old_config = [
        {
            "url": "http://127.0.0.1:8777/api/search",
            "name": "AI 智能搜题",
            "method": "post",
            "contentType": "json",
            "data": {"title": "${title}", "type": "${type}", "options": "${options}"},
            "handler": "return (res)=> [res.question, res.answer]",
        }
    ]
    new_config = generate_config()
    record("旧配置会被 OCS 判定为不支持（复现警告）", ocs_supports_image(old_config) is False)
    record("新配置满足 OCS 5.0 图片支持检测", ocs_supports_image(new_config) is True)

    data = new_config[0]["data"]
    record("data 含 images handler", isinstance(data.get("images"), dict))
    record(
        "data 含 suggestion_title/options",
        data.get("suggestion_title") == "${suggestion_title}"
        and data.get("suggestion_options") == "${suggestion_options}",
    )
    record("方法为 POST", new_config[0].get("method") == "post")


# ═══════════════ E2E 部分 ═══════════════


def e2e_checks() -> None:
    print("\n[E2E] 通过 FastAPI TestClient 真实调用 POST /api/search")
    from fastapi.testclient import TestClient

    from server import app

    title_img = make_image("Solve for x:  3x + 5 = 20")
    opt_imgs = [make_image(str(n), 200, 140) for n in (3, 5, 7, 9)]

    def norm(answer: str) -> str:
        answer = (answer or "").upper()
        m = re.search(r"(?<![A-Z])([A-Z])(?![A-Z])", answer)
        if m:
            return m.group(1)
        letters = re.findall(r"[A-Z]", answer)
        if letters and len(letters) == 1:
            return letters[0]
        return answer.strip(" .,:;，。：；")

    def call(name: str, payload: dict, expect: str) -> None:
        with TestClient(app) as client:
            resp = client.post("/api/search", json=payload)
        ok_http = resp.status_code == 200
        data = resp.json() if ok_http else {}
        answer = str(data.get("answer", ""))
        ok = ok_http and data.get("code") == 1 and norm(answer) == expect
        record(name, ok, f"HTTP {resp.status_code} code={data.get('code')} answer={answer!r}")
        if not ok and data.get("msg"):
            print("       msg:", data.get("msg")[:300])

    with TestClient(app) as client:
        health = client.get("/health")
    print(f"  health: {health.status_code} {health.json()}")

    # 1) title 图片：原始 title 里仍有坏 URL，但走 images 数组，不应去下载
    payload_title_img = {
        "title": "求 x 的值。题目图片：https://p.ananas.chaoxing.com/star3/origin/not-found.png",
        "type": "single",
        "options": "A. 3\nB. 5\nC. 7\nD. 9",
        "images": [title_img],
        "suggestion_title": "求 x 的值。题目图片： [图片1] ",
        "suggestion_options": "",
    }
    call("title 图片（base64 + [图片1]）", payload_title_img, "B")

    # 2) 纯图片选项
    payload_option_imgs = {
        "title": "Solve for x: 3x + 5 = 20",
        "type": "single",
        "options": (
            "A. https://p.ananas.chaoxing.com/star3/origin/a.png\n"
            "B. https://p.ananas.chaoxing.com/star3/origin/b.png\n"
            "C. https://p.ananas.chaoxing.com/star3/origin/c.png\n"
            "D. https://p.ananas.chaoxing.com/star3/origin/d.png"
        ),
        "images": opt_imgs,
        "suggestion_title": "Solve for x: 3x + 5 = 20",
        "suggestion_options": "A. [图片1]\nB. [图片2]\nC. [图片3]\nD. [图片4]",
    }
    call("纯图片选项（4 张 base64）", payload_option_imgs, "B")

    # 3) 混合：题干图片 + 4 张选项图片
    payload_mixed = {
        "title": "求 x 的值：https://p.ananas.chaoxing.com/star3/origin/q.png",
        "type": "single",
        "options": "A. 3\nB. 5\nC. 7\nD. 9",
        "images": [title_img, make_image("3", 200, 140), make_image("5", 200, 140),
                   make_image("7", 200, 140), make_image("9", 200, 140)],
        "suggestion_title": "求 x 的值： [图片1] ",
        "suggestion_options": "A. [图片2]\nB. [图片3]\nC. [图片4]\nD. [图片5]",
    }
    call("混合多图（5 张 base64）", payload_mixed, "B")

    # 4) 纯文字题：images 为空、suggestion 为空（OCS 不发这些字段时的正常形态）
    payload_text = {
        "title": "2 + 3 = ?",
        "type": "single",
        "options": "A. 4\nB. 5\nC. 6\nD. 7",
        "images": [],
        "suggestion_title": "",
        "suggestion_options": "",
    }
    call("纯文字回归", payload_text, "B")

    # 5) 旧版 payload（没有 images/suggestion_* 字段）兼容
    payload_old = {
        "title": "2 + 3 = ?",
        "type": "single",
        "options": "A. 4\nB. 5\nC. 6\nD. 7",
    }
    call("旧版 payload 兼容", payload_old, "B")

    # 6) 老配置把占位符替换成 undefined 字面量时不报错
    payload_undefined = {
        "title": "2 + 3 = ?",
        "type": "single",
        "options": "A. 4\nB. 5\nC. 6\nD. 7",
        "images": [],
        "suggestion_title": "undefined",
        "suggestion_options": "undefined",
    }
    call("undefined 字面量容错", payload_undefined, "B")


def main() -> int:
    print("=" * 72)
    print("OCS 5.0 图片题优化链路测试")
    print("=" * 72)
    offline_checks()
    if "--offline-only" not in sys.argv:
        e2e_checks()
    print("\n" + "=" * 72)
    if FAILS:
        print(f"总体: FAIL ❌  ({len(FAILS)} 项失败: {', '.join(FAILS)})")
        return 1
    print("总体: PASS ✅")
    return 0


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