#!/usr/bin/env python3
"""Submit scene 01 atmosphere batches to MXAI.
This script only encodes reference files for API input and downloads API outputs; it does not edit/process images locally.
"""
import base64, json, mimetypes, os, pathlib, re, time, urllib.error, urllib.request
from datetime import datetime

BASE = "https://mcp.mxai.cn"
ROOT = pathlib.Path('/Users/bot1/AI Work/02_projects/active/proj-coin-pouch-video')
OUTDIR = ROOT / 'outputs/scene-01-base/scene01-atmosphere-redo-20260522-161016'
OUTDIR.mkdir(parents=True, exist_ok=True)

PROMPT = """参考白底排版图和10张原始钱币图，生成一张 3:4 竖版场景01氛围图。

画面是一组10枚不同形态的古钱币文物原型，错落摆放在深色暖调桌面/展台上。保留10枚钱币的形态差异：桥形灰金属钱、贝形金色钱、布币/铲形钱、青铜绿锈铲形钱、圆形方孔钱、黑色椭圆刻纹钱、金色裂纹圆片等都要可辨识，不能统一成同一种圆形铜钱。

构图参考白底排版图的数量、位置和错落关系。10枚钱币需要可数、可辨，至少6枚完整或大部分可见，其余可以轻微搭边、压住、半遮挡，但轮廓、孔洞、颜色或纹样仍能辨认。允许自然叠放和半遮挡，画面要丰富，但不要堆成一团。

氛围为文博奇幻、微缩世界、温暖神秘。深棕黑色桌面，细腻纹理和轻微尘埃颗粒。一束柔和暖金光从左上方斜照，扫过钱币边缘和纹样，形成低调金色高光。背景虚化，浅景深，微距摄影质感，整体有高级产品广告感和电影感。

镜头为微距近景，轻微斜俯拍视角，像视频开场第一帧。画面中心有视觉焦点，但不要只突出一枚钱币，需要让10枚钱币作为一个群体被看见。

避免：文字、水印、边框、拼贴感、白底素材板、钱币堆成一团、10枚变成同一种圆形铜钱、明显悬浮、主体模糊、廉价卡通感。"""

REFS = [
    ROOT / 'outputs/scene-01-base/scene01-ten-coin-staggered-layout-feathered-20260522-135944.png',
] + sorted((ROOT / 'assets/reference-coins/raw').glob('*.jpg'))


def redact(s):
    s = str(s)
    s = re.sub(r'nb_[A-Za-z0-9_\-]+', 'nb_********', s)
    s = re.sub(r'(Authorization: Bearer )[A-Za-z0-9_\-\.]+', r'\1********', s)
    return s


def data_url(path):
    mime = mimetypes.guess_type(str(path))[0] or 'image/png'
    b64 = base64.b64encode(path.read_bytes()).decode('ascii')
    return f'data:{mime};base64,{b64}'


def request_json(method, path, key, body=None, timeout=120):
    headers = {'Content-Type':'application/json; charset=utf-8', 'Authorization':'Bearer '+key}
    data = None if body is None else json.dumps(body, ensure_ascii=False).encode('utf-8')
    req = urllib.request.Request(BASE+path, data=data, method=method, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.loads(r.read().decode('utf-8', errors='replace'))
    except urllib.error.HTTPError as e:
        raw = e.read().decode('utf-8', errors='replace')
        raise RuntimeError(f'HTTP {e.code} {path}: {redact(raw)}')


def obj_data(resp):
    if isinstance(resp, dict) and isinstance(resp.get('data'), dict):
        return resp['data']
    return resp if isinstance(resp, dict) else {}


def get_serial(resp):
    d = obj_data(resp)
    return d.get('serial_no') or d.get('serialNo') or d.get('serial') or resp.get('serial_no')


def download(url, dest):
    req = urllib.request.Request(url, headers={'User-Agent':'HermesVideo/1.0'})
    with urllib.request.urlopen(req, timeout=180) as r:
        dest.write_bytes(r.read())


def run_batch(label, model, input_images, key):
    body = {
        'prompt': PROMPT,
        'model': model,
        'aspect_ratio': '3:4',
        'resolution': '1K',
        'quality': 'high',
        'count': 3,
        'input_images': input_images,
    }
    resp = request_json('POST', '/mcp/api/generate/image', key, body, timeout=180)
    serial = get_serial(resp)
    if not serial:
        raise RuntimeError(f'No serial for {label}: {redact(resp)}')
    print(json.dumps({'event':'submitted','label':label,'model':model,'serial_no':serial}, ensure_ascii=False), flush=True)
    rec = {'label': label, 'model': model, 'serial_no': serial, 'status': 'submitted', 'files': []}
    deadline = time.time() + 900
    while time.time() < deadline:
        task = obj_data(request_json('GET', f'/mcp/api/task/{serial}', key, timeout=60))
        status = str(task.get('status'))
        print(json.dumps({'event':'poll','label':label,'serial_no':serial,'status':status}, ensure_ascii=False), flush=True)
        if status == '2':
            urls = task.get('image_urls') or task.get('images') or task.get('urls') or []
            if isinstance(urls, str):
                urls = [urls]
            for i, url in enumerate(urls, 1):
                ext = '.png'
                m = re.search(r'\.(png|jpg|jpeg|webp)(?:\?|$)', url, re.I)
                if m:
                    ext = '.' + m.group(1).lower().replace('jpeg','jpg')
                dest = OUTDIR / f'{label}-{i:02d}{ext}'
                download(url, dest)
                rec['files'].append(str(dest))
            rec['status'] = 'completed'
            return rec
        if status in ('3','4'):
            rec['status'] = 'failed'
            rec['fail_msg'] = task.get('fail_msg') or task.get('message') or redact(task)
            return rec
        time.sleep(5)
    rec['status'] = 'timeout'
    return rec


def main():
    key = os.environ.get('MX_AI_API_KEY')
    if not key:
        raise SystemExit('MX_AI_API_KEY missing')
    input_images = [data_url(p) for p in REFS]
    results = []
    # Batch 1: Nano Banana / Nano series
    results.append(run_batch('nanobanana', 'nano-2.0-pro', input_images, key))
    print(json.dumps({'event':'batch_interval_sleep','seconds':30}, ensure_ascii=False), flush=True)
    time.sleep(30)
    # Batch 2: GPT Image 2
    results.append(run_batch('gpt2image', 'gpt-image-2', input_images, key))
    summary = {'ok': True, 'outdir': str(OUTDIR), 'reference_count': len(REFS), 'results': results}
    (OUTDIR / 'mxai-generation-summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')
    print('FINAL_JSON_START')
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    print('FINAL_JSON_END')

if __name__ == '__main__':
    main()