#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片相关性守卫 — 检测「配图与标题/摘要不匹配」的条目并自动重生成

背景（2026-09-21 用户纠正）：
  真正的问题不是 AI 水印，而是**配图内容与标题不符**。
  根因：search_list_op --img-only 的选图逻辑只按「尺寸 + 宽高比」打分挑页面最大横图，
  不判断内容相关性 → 常抓到文章页的无关配图（相关推荐缩略图、游戏插画、站点装饰图、
  其他新闻的封面等）。实例：标题「"卧薪尝胆"棒棒糖走红」抓到《梦幻西游》游戏插画。

判定标准（视觉模型逐条比对标题）：
  ✓ 图与标题为「是」= 主体/事件/对象一致 → 保留
  ✗ 「否」或「部分」→ 判定不匹配 → 豆包按标题重生成

附带处理：豆包 API 生成的图右下角自带「AI生成」水印（底部 60px 内），
        重生成后自动裁除并裁到 16:9、缩放 750x422。

用法：
  python3.12 scripts/image_relevance_guard.py --scan <json> <images_dir>
  python3.12 scripts/image_relevance_guard.py --fix  <json> <images_dir> [--only brand_hotspots_0,social_hotspots_3]

同步：修复后自动把 thumb 保持原路径（文件名不变），无需改 JSON；如需改 thumb 路径会打印提示。
"""
import argparse
import base64
import json
import os
import re
import urllib.request

from PIL import Image

ARK_KEY = "ark-563a4210-d804-47fa-a37a-40962192119b-ce6db"
GEN_URL = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
CHAT_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
GEN_MODEL = "doubao-seedream-4-0-250828"
VISION_MODEL = "doubao-seed-2-1-pro-260628"
SEC = ['brand_hotspots', 'vehicle_hotspots', 'social_hotspots']

WATERMARK_CUT = 60
TARGET_W = 750
PEOPLE_BAN = "，无人物，不出现人脸、人体、手部、剪影"
# 2026-09-22：ark 视觉接口偶发 socket 挂起（单次最长 120s，重试 2 次 = 4 分钟/张）。
# 改为环境变量可调，批量扫描时用 RELEVANCE_TIMEOUT=45 RELEVANCE_TRY=1 快速失败。
VISION_TIMEOUT = int(os.environ.get('RELEVANCE_TIMEOUT', '120'))
VISION_TRY = int(os.environ.get('RELEVANCE_TRY', '2'))

# 2026-09-22 用户规则 A：**内容里抓取的（真实）新闻照片含人物不做处理**。
# 「禁人物」只约束豆包生成图的 prompt（PEOPLE_BAN），不用于否决已下载的真实照片。
# 因此默认忽略判定项 3)，只按「图文是否匹配」决定是否重生成。
# 需要恢复旧行为（含人物也重生成）时加 --flag-people。
ALLOW_PEOPLE = True


def vision_ask(path, question, max_try=None):
    max_try = VISION_TRY if max_try is None else max_try
    last = ''
    for _ in range(max_try):
        try:
            b64 = base64.b64encode(open(path, 'rb').read()).decode()
            body = json.dumps({"model": VISION_MODEL, "messages": [{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64}},
                {"type": "text", "text": question}]}]}).encode()
            req = urllib.request.Request(CHAT_URL, data=body, headers={
                "Content-Type": "application/json", "Authorization": f"Bearer {ARK_KEY}"})
            return json.loads(urllib.request.urlopen(req, timeout=VISION_TIMEOUT).read())["choices"][0]["message"]["content"].strip().replace("\n", " ")
        except Exception as e:
            last = str(e)[:80]
    return f"__ERR__{last}"


def relevance_ask(path, title, summary=''):
    q = (f"这是一条新闻的配图。新闻标题：「{title}」。摘要：「{(summary or '')[:60]}」。\n"
         "请判断配图与这条新闻是否匹配，只回答三点，每点一行，30字内：\n"
         "1)是否匹配？(是/部分/否)\n"
         "2)画面主体是什么？\n"
         "3)是否出现人脸、人体、手部或人物剪影？(是/否)")
    return vision_ask(path, q)


def parse_relevance(raw):
    if raw.startswith('__ERR__'):
        return None, '视觉接口失败', raw
    m = re.search(r'1\)?\s*[:：]?\s*(是|部分|否)', raw)
    verdict = m.group(1) if m else '?'
    m3 = re.search(r'3\)?\s*[:：]?\s*(是|否)', raw)
    people = (m3.group(1) == '是') if m3 else False
    problems = []
    if verdict in ('否', '部分', '?'):
        problems.append(f'图文不匹配({verdict})')
    if people and not ALLOW_PEOPLE:
        problems.append('含人物')
    return (len(problems) == 0), ','.join(problems), raw


def gen_image(prompt):
    body = json.dumps({"model": GEN_MODEL, "prompt": prompt, "size": "1280x720",
                       "n": 1, "response_format": "url"}).encode()
    req = urllib.request.Request(GEN_URL, data=body, headers={
        "Content-Type": "application/json", "Authorization": f"Bearer {ARK_KEY}"})
    u = json.loads(urllib.request.urlopen(req, timeout=180).read())["data"][0]["url"]
    return urllib.request.urlopen(urllib.request.Request(
        u, headers={"User-Agent": "Mozilla/5.0"}), timeout=120).read()


def strip_watermark(im):
    w, h = im.size
    if h > WATERMARK_CUT + 100:
        im = im.crop((0, 0, w, h - WATERMARK_CUT))
    w2, h2 = im.size
    tw = int(round(h2 * 16 / 9))
    if tw < w2:
        left = (w2 - tw) // 2
        im = im.crop((left, 0, left + tw, h2))
    return im.resize((TARGET_W, int(round(TARGET_W * 9 / 16))), Image.LANCZOS)


def clean_for_prompt(s, maxlen=60):
    """去热度值、去人名（避免豆包画人物肖像）"""
    s = re.sub(r'\s+\d+(?:\.\d+)?\s*万\s*(?:热度)?\s*$', '', s or '')
    s = re.sub(r'[\u4e00-\u9fff]{2,4}(?:副部长|部长|总裁|CEO|总经理|董事长|总工程师|教授|博士)', '', s)
    s = re.sub(r'[A-Za-z][A-Za-z\s.]{0,20}(?=：|:)', '', s)
    return s.strip(' ，。:：')[:maxlen]


def build_prompts(title, summary):
    t = clean_for_prompt(title)
    s = clean_for_prompt(summary, 50)
    return [
        (f"配图：{t}。{s}。新闻纪实配图，写实摄影，摄影棚布光，单反真实拍摄，"
         f"画面主体紧扣新闻主题，横向构图，16:9，高清" + PEOPLE_BAN),
        (f"配图：{t}。围绕该新闻主题的实物静物特写，浅景深，商业摄影，"
         f"真实照片质感，横向构图，16:9，高清" + PEOPLE_BAN),
        (f"配图：{t}。该新闻场景的环境/场景空镜，无人物，写实摄影，"
         f"横向构图，16:9，高清" + PEOPLE_BAN),
    ]


def _sections(data):
    """兼容两种顶层格式（2026-09-22 修复）：
       dict 顶层 = 最终文件 json/{MMDD}data.json
       list 顶层 = yes_data / origin_data（元素形如 {'name': 板块名, 'list': [...]}）
    """
    if isinstance(data, dict):
        return data
    out = {}
    for seg in data:
        if isinstance(seg, dict) and 'name' in seg:
            out[seg['name']] = seg.get('list', [])
    return out


def iter_items(data):
    secs = _sections(data)
    for k in SEC:
        for i, it in enumerate(secs.get(k, [])):
            th = it.get('thumb', '') or ''
            if th.startswith('/'):
                yield k, i, os.path.basename(th), it


def scan(json_path, img_dir, only=None):
    data = json.load(open(json_path))
    bad = []
    for k, i, fname, it in iter_items(data):
        if only and fname not in only:
            continue
        p = os.path.join(img_dir, fname)
        if not os.path.exists(p):
            print(f"  ❌ {k}[{i}] {fname} 文件不存在")
            bad.append(fname)
            continue
        ok, why, raw = parse_relevance(relevance_ask(p, it.get('title', ''), it.get('summary', '')))
        if ok:
            print(f"  ✅ {k}[{i}] {fname:26s} | {raw[:80]}")
        else:
            print(f"  ⚠️ {k}[{i}] {fname:26s} → {why}")
            print(f"       标题: {it.get('title','')}")
            print(f"       画面: {raw[:110]}")
            bad.append(fname)
    return bad


def fix(json_path, img_dir, only=None):
    data = json.load(open(json_path))
    fixed = 0
    for k, i, fname, it in iter_items(data):
        if only and fname not in only:
            continue
        p = os.path.join(img_dir, fname)
        ok, why, raw = parse_relevance(relevance_ask(p, it.get('title', ''), it.get('summary', '')))
        if ok:
            print(f"  ✅ {k}[{i}] {fname:26s} 保留")
            continue
        print(f"  ⚠️ {k}[{i}] {fname} → {why} → 按标题重生成")
        print(f"       标题: {it.get('title','')}")
        # 2026-09-22 修复：候选是直接写进最终路径再校验的，若全部不合规，
        #   盘上会残留「最后一张不合规候选」而日志只说「需人工处理」→ 原图被静默替换。
        #   现提前备份原图，全部候选失败时还原。
        _orig = open(p, 'rb').read() if os.path.exists(p) else None
        done = False
        for ci, pr in enumerate(build_prompts(it.get('title', ''), it.get('summary', '')), 1):
            if done:
                break
            try:
                tmp = f"/tmp/rg_{fname}_{ci}.jpg"
                open(tmp, 'wb').write(gen_image(pr))
                im = Image.open(tmp).convert('RGB')
                strip_watermark(im).save(p, 'JPEG', quality=92)
                ok2, why2, raw2 = parse_relevance(relevance_ask(p, it.get('title', ''), it.get('summary', '')))
                if ok2:
                    o = Image.open(p)
                    print(f"       ✅ 候选{ci} 通过 → {o.format} {o.size} {os.path.getsize(p)//1024}KB")
                    print(f"          {raw2[:100]}")
                    fixed += 1
                    done = True
                else:
                    print(f"       ⚠️ 候选{ci} 不合规({why2}) | {raw2[:70]}")
            except Exception as e:
                print(f"       ❌ 候选{ci} 异常: {str(e)[:90]}")
        if not done:
            if _orig is not None:
                open(p, 'wb').write(_orig)
                print(f"       ↩️ 全部候选不合规 → 已还原原图（未被静默替换）")
            print(f"       ⚠️ {fname} 需人工处理")
    print(f"\n合计修复 {fixed} 张")


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--scan', nargs=2, metavar=('JSON', 'IMAGES_DIR'))
    ap.add_argument('--fix', nargs=2, metavar=('JSON', 'IMAGES_DIR'))
    ap.add_argument('--only', default='',
                    help='只处理指定图片，逗号分隔。可写不带扩展名的名字（如 brand_hotspots_0），'
                         '脚本会自动补 .jpg/.png/.jpeg（2026-09-22 加，此前必须带 .jpg 否则静默 0 结果）')
    ap.add_argument('--flag-people', action='store_true',
                    help='恢复旧行为：真实照片含人物也判为需重生成（默认按用户规则A忽略人物）')
    a = ap.parse_args()
    if a.flag_people:
        ALLOW_PEOPLE = False
        print('⚠️ --flag-people：真实照片含人物也判为需重生成')
    only = set(x.strip() for x in a.only.split(',') if x.strip()) or None
    if only:
        # 容错归一化：不带扩展名时自动补全（内部用 os.path.basename 比对）
        norm = set(only)
        for x in only:
            if '.' not in os.path.basename(x):
                for ext in ('.jpg', '.jpeg', '.png'):
                    norm.add(x + ext)
        only = norm
    if not only and a.only:
        print('⚠️ --only 解析为空，将处理全部条目')
    if a.scan:
        print(f"=== 图文相关性扫描: {a.scan[0]} ===")
        b = scan(a.scan[0], a.scan[1], only)
        print(f"\n不匹配: {len(b)} 张 -> {b}")
    elif a.fix:
        print(f"=== 图文相关性修复: {a.fix[0]} ===")
        fix(a.fix[0], a.fix[1], only)
    else:
        ap.print_help()