#!/usr/bin/env python3.12
"""批量豆包重生成小尺寸/PNG伪装图片 (brand 0-4, vehicle 0-4)"""
import json, urllib.request, os, re, sys, time

API_KEY = "ark-563a4210-d804-47fa-a37a-40962192119b-ce6db"
API_URL = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
MODEL = "doubao-seedream-4-0-250828"
IMG_DIR = "/data/news/images/0813/"
JSON_PATH = "/data/news/json/yes_data/0812data.json"

def gen_image(prompt, fname):
    payload = json.dumps({
        "model": MODEL,
        "prompt": prompt,
        "size": "1280x720",
        "n": 1,
        "response_format": "url"
    }).encode()
    req = urllib.request.Request(API_URL, data=payload,
        headers={'Authorization': f'Bearer {API_KEY}', 'Content-Type': 'application/json'})
    resp = urllib.request.urlopen(req, timeout=120)
    result = json.loads(resp.read().decode('utf-8'))
    img_url = result.get('data', [{}])[0].get('url', '')
    if not img_url:
        raise Exception("豆包未返回图片URL")
    local_path = os.path.join(IMG_DIR, fname)
    img_req = urllib.request.Request(img_url, headers={'User-Agent': 'Mozilla/5.0'})
    img_resp = urllib.request.urlopen(img_req, timeout=60)
    with open(local_path, 'wb') as f:
        f.write(img_resp.read())
    return local_path

def clean_title(t):
    return re.sub(r'\s*\d[\d,.]*\s*万(?:热度)?\s*$', '', t)

def clean_summary(s):
    if any(kw in s for kw in ['环境异常', '验证码', '反爬', '页面加载']):
        return ''
    return s

d = json.load(open(JSON_PATH))
items = d if isinstance(d, list) else [d]

targets = []
for s in items:
    name = s.get('name')
    if name == 'brand_hotspots':
        for idx in [0,1,2,3,4]:
            if idx < len(s['list']):
                targets.append((name, idx))
    elif name == 'vehicle_hotspots':
        for idx in [0,1,2,3,4]:
            if idx < len(s['list']):
                targets.append((name, idx))

for name, idx in targets:
    fname = f'{name}_{idx}.jpg'
    section = next(s for s in items if s.get('name') == name)
    it = section['list'][idx]
    title = clean_title(it.get('title',''))
    summary = clean_summary(it.get('summary',''))
    prompt = f"配图：{title}，{summary[:60]}，新闻配图风格，写实，横向构图，16:9，高清"
    print(f'[{name}][{idx}] 生成中... {title[:35]}')
    try:
        path = gen_image(prompt, fname)
        size = os.path.getsize(path)
        print(f'  ✅ {fname} ({size//1024}KB)')
    except Exception as e:
        print(f'  ❌ 失败: {e}')
    time.sleep(1)

print('\n✅ 批量重生成完成')