"""Probe: 方舟"文档理解"能不能直读扫描件 PDF(路线 A 验证,不动线上代码)。

跑法: .venv/Scripts/python.exe scripts/probe_ark_doc.py [真实PDF路径]
依赖 .env 里 ARK_API_KEY。**会真调豆包 seed-2.0-lite,产生 < ¥0.05 费用**。

验证点:
  1. 合成一份 3 页"扫描件"PDF(PIL 纯图页,无文本层,每页埋魔术串 + 中文段落/表格)
  2. markitdown 对它抽不出文字(证明现状确实是死路)
  3. chat/completions 的 file 内容块用哪种 JSON 形状能被接受(候选格式挨个试,
     400 报错会带字段名,本身就是探针产出)
  4. OCR 保真度:三页魔术串是否全命中(页覆盖)、中文/表格内容是否读出
  5. usage tokens → 每页成本口径

传真实 PDF 路径时跳过合成,直接测该文件(页数/体积上限、真实扫描件质量)。
"""
from __future__ import annotations

import base64
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

# Windows 控制台默认 GBK,打印中文/特殊符号会崩 → 强制 stdout UTF-8
try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")  # type: ignore[attr-defined]
except Exception:
    pass

# 读 .env(同 smoke_look_at_image)
env_file = ROOT / ".env"
if env_file.exists():
    for line in env_file.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, _, v = line.partition("=")
        os.environ.setdefault(k.strip(), v.strip())

from PIL import Image, ImageDraw, ImageFont

from core.ark_client import ArkClient, ArkConfig, ArkError

def _magic(i: int) -> str:
    """第 i 页(0 起)的魔术串,确定性可复算。"""
    return f"ZCBOT-DOC-P{i + 1}-{7391 + i * 613}"

# 每页正文:标题 + 段落 + 第 2 页一个小表格,模拟真实扫描标准的版面
PAGE_LINES = [
    [
        "水泥胶砂强度检验方法(模拟扫描件 第1页)",
        "本方法规定了水泥胶砂抗压强度与抗折强度的测定步骤。",
        "试验室温度应保持在 20 ± 2 摄氏度,相对湿度不低于 50%。",
        "校验码: {magic}",
    ],
    [
        "第2页 配合比与龄期",
        "胶砂配比  水泥 450g  标准砂 1350g  水 225g",
        "龄期(d)   3      7      28",
        "抗压(MPa) 22.5   35.0   52.5",
        "校验码: {magic}",
    ],
    [
        "第3页 结果处理",
        "以三条试体抗折结果的算术平均值作为试验结果。",
        "当三个值中有超出平均值 ±10% 时,应剔除后重新计算。",
        "校验码: {magic}",
    ],
]

_FONT_CANDIDATES = [
    r"C:\Windows\Fonts\msyh.ttc",
    r"C:\Windows\Fonts\simhei.ttf",
    r"C:\Windows\Fonts\simsun.ttc",
]


def _load_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    for p in _FONT_CANDIDATES:
        if Path(p).exists():
            return ImageFont.truetype(p, size)
    return ImageFont.load_default()


def make_scanned_pdf(dest: Path, n_pages: int = 3) -> None:
    """N 页纯图 PDF(A4 150dpi,无文本层)—— 模拟扫描件。

    前 3 页用真实版面(段落 + 表格),之后的页是短页(标题 + 魔术串),
    专测页数上限 / 页覆盖。1-bit 模式控体积(黑字白底扫描件本就近似双色)。
    """
    font = _load_font(36)
    pages = []
    for i in range(n_pages):
        img = Image.new("RGB", (1240, 1754), (255, 255, 255))
        d = ImageDraw.Draw(img)
        lines = (PAGE_LINES[i] if i < len(PAGE_LINES)
                 else [f"第{i + 1}页 附录条款", "本页为附录占位内容。", "校验码: {magic}"])
        y = 120
        for line in lines:
            d.text((100, y), line.format(magic=_magic(i)), fill=(0, 0, 0), font=font)
            y += 90
        # 1-bit 模式绕开本环境 PIL 缺 JPEG 编码器的问题(RGB 页会走 DCT/JPEG),
        # 且体积最小(P 模式实测一页 4MB+,多页测试撑爆请求)
        pages.append(img.convert("1", dither=Image.Dither.NONE))
    # resolution=150:页物理尺寸=px/150 英寸 ≈ A4。缺省 72dpi 会把页标成 A4 两倍大,
    # 方舟按固定 dpi 栅格化 PDF 页,超尺寸页撞"单页 3600 万像素"上限(probe 实测报错)
    pages[0].save(dest, save_all=True, append_images=pages[1:], resolution=150.0)


def check_markitdown_dead_end(pdf: Path) -> None:
    """现状对照:markitdown 对纯图 PDF 应抽不出正文。"""
    exe = ROOT / ".venv" / "Scripts" / "markitdown.exe"
    cmd = [str(exe) if exe.exists() else "markitdown", str(pdf)]
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=120,
                           encoding="utf-8", errors="replace")
        text = (r.stdout or "").strip()
        print(f"[markitdown] exit={r.returncode} 抽出正文 {len(text)} 字符"
              + (f" → 非空?! 前 200 字: {text[:200]!r}" if text else " → 空(证实扫描件死路)"))
    except Exception as e:
        print(f"[markitdown] 跑不了({type(e).__name__}: {e}),跳过对照")


# file 内容块候选形状:方舟文档说 file_id/file_data/file_url 三选一,但块的外层
# JSON 没抓到 → 挨个试,400 报错信息(带字段名)也是探针产出
def _candidate_blocks(b64: str, filename: str) -> list[tuple[str, dict]]:
    data_url = f"data:application/pdf;base64,{b64}"
    return [
        ("openai_file_data_url",
         {"type": "file", "file": {"filename": filename, "file_data": data_url}}),
        ("openai_file_raw_b64",
         {"type": "file", "file": {"filename": filename, "file_data": b64}}),
        ("ark_file_url_data",
         {"type": "file_url", "file_url": {"url": data_url}}),
    ]


QUESTION = (
    "这是一份多页 PDF 文档。请逐页把其中的文字完整 OCR 出来,"
    "每页以「== 第N页 ==」开头,保留表格数据与换行,不要总结不要遗漏。"
)

# 多页模式只要校验码清单:把"页覆盖上限"和"输出 token 上限"两个变量拆开测
QUESTION_MAGIC_ONLY = (
    "这是一份多页 PDF 文档,每页都有一行「校验码: ZCBOT-DOC-...」。"
    "请按页序把每页的校验码逐行列出(格式:第N页 <校验码>),只要校验码,别的不用输出。"
)


def main() -> int:
    cfg = ArkConfig.load()
    if cfg is None:
        print("[SKIP] ARK_API_KEY 未设(或 doubao.yaml 缺失)")
        return 0
    vcfg = (cfg.raw.get("vision") or {}).get("seed_2_lite") or {}
    model_id = vcfg.get("model_id", "doubao-seed-2-0-lite-260428")
    print(f"[setup] model={model_id} base={cfg.base_url}")

    n_pages = 3
    if len(sys.argv) > 2 and sys.argv[1] == "--pages":
        n_pages = int(sys.argv[2])
        pdf = None
    elif len(sys.argv) > 1:
        pdf = Path(sys.argv[1])
    else:
        pdf = None

    if pdf is not None:
        synthetic = False
        print(f"[setup] 使用真实 PDF: {pdf}")
    else:
        pdf = Path(tempfile.mkdtemp(prefix="zcbot_probe_")) / "scanned_probe.pdf"
        make_scanned_pdf(pdf, n_pages)
        synthetic = True
        print(f"[setup] 合成扫描件 PDF: {pdf}"
              f"({pdf.stat().st_size} bytes, {n_pages} 页, 无文本层)")

    check_markitdown_dead_end(pdf)

    b64 = base64.b64encode(pdf.read_bytes()).decode()
    print(f"[setup] base64 体积 {len(b64) / 1024:.0f} KB")

    question = QUESTION_MAGIC_ONLY if (synthetic and n_pages > 5) else QUESTION
    resp = None
    accepted = None
    for name, block in _candidate_blocks(b64, pdf.name):
        body = {
            "model": model_id,
            "messages": [{
                "role": "user",
                "content": [{"type": "text", "text": question}, block],
            }],
        }
        print(f"[try] 格式 {name} ...")
        try:
            with ArkClient(cfg, timeout_s=300) as client:
                resp = client.post_json("/chat/completions", body, timeout_s=300)
            accepted = name
            print(f"[OK] 格式 {name} 被接受")
            break
        except ArkError as e:
            print(f"[reject] {name}: {e}")

    if resp is None:
        print("\n[FAIL] 所有候选格式都被拒 —— 看上面报错定位正确字段名,"
              "或该模型版本不支持 file 输入(需查文档理解模型列表)")
        return 2

    content = ((resp.get("choices") or [{}])[0].get("message") or {}).get("content") or ""
    if isinstance(content, list):
        content = "\n".join(c.get("text", "") for c in content if isinstance(c, dict))
    usage = resp.get("usage") or {}
    tin = int(usage.get("prompt_tokens", 0) or 0)
    tout = int(usage.get("completion_tokens", 0) or 0)
    cost = (tin * float(vcfg.get("price_cny_per_mtoken_input", 0.6))
            + tout * float(vcfg.get("price_cny_per_mtoken_output", 3.6))) / 1e6

    print(f"\n[usage] tokens={tin}+{tout} cost≈¥{cost:.4f} accepted_format={accepted}")
    print(f"[response]\n{content}\n")

    if synthetic:
        flat = content.replace(" ", "").replace("-", "")
        magics = [_magic(i) for i in range(n_pages)]
        hits = [m for m in magics if m.replace("-", "") in flat]
        missed = [m for m in magics if m not in hits]
        print(f"[verify] 魔术串命中 {len(hits)}/{n_pages}"
              + (f" 漏: {missed[:5]}{'...' if len(missed) > 5 else ''}" if missed else ""))
        if n_pages <= 5:
            cn_hit = "标准砂" in content and "抗压" in content
            print(f"[verify] 中文表格关键词(标准砂/抗压)命中: {cn_hit}")
        else:
            cn_hit = True
        if not missed and cn_hit:
            print(f"\n[PASS] {n_pages} 页全覆盖:file 输入 + 扫描件 OCR 验证通过")
            return 0
        print("\n[WARN] 部分未命中 —— 人工核对上面 response 判断保真度")
        return 1
    print("[DONE] 真实 PDF 模式:人工核对上面 response")
    return 0


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