#!/usr/bin/env python3
from __future__ import annotations
import base64, datetime as dt, json, mimetypes, os, re, subprocess, sys, time, urllib.request, urllib.error
from pathlib import Path
BASE='https://mcp.mxai.cn'
MODEL='gpt-image-2'; ASPECT_RATIO='9:16'; RESOLUTION='1K'; QUALITY='high'
PROJECT=Path('/Users/bot1/Volumes/root_for_ai/AI工作区/国博古钱币盲盒套装_商品主视频_20260625')
SRC_BASE=Path('/Users/bot1/Volumes/root_for_ai/public-assets/product-archive-library/archive/2026/PA-2026-0006__国博古钱币盲盒套装/01_拍摄原片')
OUT_FRAMES=PROJECT/'02_storyboard_frames_mxai'; OUT_PROMPTS=PROJECT/'04_prompts'; REF_CACHE=PROJECT/'01_refs_mxai_cache'
for d in [OUT_FRAMES, OUT_PROMPTS, REF_CACHE]: d.mkdir(parents=True, exist_ok=True)
MANIFEST=OUT_PROMPTS/'mxai_storyboard_generation_manifest.json'
RUN_ID='20260625_160244'
KEY=os.environ.get('MX_AI_API_KEY')
if not KEY: raise SystemExit('MX_AI_API_KEY missing')

def redact(text, key=KEY):
    s=str(text)
    if key: s=s.replace(key,'********')
    s=re.sub(r'nb_[A-Za-z0-9_\-]+','nb_********',s)
    return s

def request_json(method,path,body=None,timeout=180):
    headers={'Content-Type':'application/json; charset=utf-8','Authorization':'Bearer '+KEY,'User-Agent':'HermesVideo/1.0'}
    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 resp:
            return json.loads(resp.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 nested(resp):
    return resp.get('data') if isinstance(resp,dict) and isinstance(resp.get('data'),dict) else (resp if isinstance(resp,dict) else {})

def serial_from(resp):
    obj=nested(resp)
    for k in ('serial_no','serialNo','serial'):
        if obj.get(k): return obj[k]
    for k in ('serial_no','serialNo','serial'):
        if isinstance(resp,dict) and resp.get(k): return resp[k]

def poll(serial, timeout_s=900):
    deadline=time.time()+timeout_s; last=None
    while time.time()<deadline:
        obj=nested(request_json('GET',f'/mcp/api/task/{serial}',None,90))
        last=obj; status=str(obj.get('status'))
        print(json.dumps({'event':'poll','serial_no':serial,'status':obj.get('status'),'message':obj.get('message') or obj.get('status_text')},ensure_ascii=False),flush=True)
        if status=='2': return obj
        if status in {'3','4'}: raise RuntimeError('MXAI task failed: '+redact(obj.get('fail_msg') or obj.get('message') or obj))
        time.sleep(5)
    raise TimeoutError('timeout '+serial+' last='+redact(last))

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

def dims(path):
    out=subprocess.check_output(['sips','-g','pixelWidth','-g','pixelHeight',str(path)],text=True,stderr=subprocess.STDOUT)
    return int(re.search(r'pixelWidth: (\d+)',out).group(1)), int(re.search(r'pixelHeight: (\d+)',out).group(1))

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

def make_ref(src):
    ref=REF_CACHE/f'{src.stem}__ref1800.jpg'
    if not ref.exists() or ref.stat().st_size==0:
        subprocess.check_call(['sips','-Z','1800','-s','format','jpeg','-s','formatOptions','88',str(src),'--out',str(ref)],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL)
    return ref

STYLE=("商品主视频分镜图，9:16竖版单镜头画面。基于参考图中的真实产品生成：保持毛绒挂件的正面造型、颜色、刺绣纹样、圆润轮廓、薄挂件比例和毛绒纤维质感。"
"所有毛绒尽量正面对镜头，像薄薄的毛绒挂件/正面纸偶小人轻轻活动；动作只表现为正面轻弹、轻晃、点头、靠近、轻碰或轻柔漂浮。"
"不要生成厚重3D公仔，不要侧面、背面、转身、奔跑、翻滚、真实开盒动作、复杂肢体拖拽。"
"背景为暖白奶油色微型展台，淡金色古钱币光点，干净电商产品宣传片质感。"
"保留产品自身已有纹样和刺绣细节；不要新增文字、额外logo、水印、标签、边框。")
ref_hidden=SRC_BASE/'20260625-140632.jpg'; ref_group=SRC_BASE/'20260625-140637.jpg'
shot5={'index':5,'id':'shot05_hidden_conductor','name':'镜头5_海晏河清尊隐藏款C位登场','ref':ref_hidden,'serial_no':'2070057454288244736','prompt':STYLE+'镜头5：海晏河清尊隐藏款C位登场。以参考图中的浅青灰色海晏河清尊原型毛绒挂件为唯一核心主体，正面位于画面中央，略高、略大，像小指挥家一样温柔抬起小手或轻轻点头；周围用淡金色光晕和少量柔和光点表现其他毛绒正在环绕，但不要把主体变成瓷器或厚重3D模型。'}
shot6={'index':6,'id':'shot06_final_group_c_position','name':'镜头6_最终全员商品主视觉','ref':ref_group,'prompt':STYLE+'镜头6：最终全员商品主视觉。七只毛绒全部以正面合影出现，浅青灰色海晏河清尊隐藏款站在正中央C位，略高于其他角色；其他六只分布在左右和下方，像用一圈柔和金色光带托起C位主角，形成完整套装合影。画面干净、可爱、适合电商商品主视频结尾定格。'}

def load_manifest():
    return json.loads(MANIFEST.read_text(encoding='utf-8')) if MANIFEST.exists() else {'run_id':RUN_ID,'project':str(PROJECT),'model':MODEL,'aspect_ratio':ASPECT_RATIO,'resolution':RESOLUTION,'quality':QUALITY,'count_per_shot':1,'shots':[]}

def save_manifest(m):
    m['completed']=sum(1 for s in m['shots'] if s.get('status')=='completed')
    m['failed']=sum(1 for s in m['shots'] if s.get('status')=='failed')
    MANIFEST.write_text(json.dumps(m,ensure_ascii=False,indent=2),encoding='utf-8')

def append_or_replace(m, rec):
    m['shots']=[s for s in m['shots'] if s.get('index')!=rec.get('index')]
    m['shots'].append(rec); m['shots'].sort(key=lambda x:x.get('index',99)); save_manifest(m)

def save_done(shot, done, serial):
    urls=done.get('image_urls') or done.get('images') or done.get('urls') or []
    if isinstance(urls,str): urls=[urls]
    if not urls: raise RuntimeError('completed task has no image url')
    files=[]
    for j,url in enumerate(urls,1):
        ext='.png'; mm=re.search(r'\.(png|jpg|jpeg|webp)(?:\?|$)',url,re.I)
        if mm: ext='.'+mm.group(1).lower().replace('jpeg','jpg')
        dest=OUT_FRAMES/f"{shot['index']:02d}_{shot['id']}_{RUN_ID}_{j}{ext}"
        download(url,dest); w,h=dims(dest)
        print(json.dumps({'event':'saved','index':shot['index'],'file':str(dest),'width':w,'height':h},ensure_ascii=False),flush=True)
        files.append({'local_file':str(dest),'width':w,'height':h,'size_bytes':dest.stat().st_size})
    rec={'index':shot['index'],'id':shot['id'],'name':shot['name'],'source_ref':str(shot['ref']),'reference_cache_file':str(make_ref(shot['ref'])),'prompt':shot['prompt'],'model':MODEL,'aspect_ratio':ASPECT_RATIO,'resolution':RESOLUTION,'quality':QUALITY,'count':1,'serial_no':serial,'status':'completed','files':files,'completed_at':dt.datetime.now().isoformat(timespec='seconds')}
    m=load_manifest(); append_or_replace(m,rec)

# poll and save shot5
print(json.dumps({'event':'continue_existing','index':5,'serial_no':shot5['serial_no']},ensure_ascii=False),flush=True)
done5=poll(shot5['serial_no'],900)
save_done(shot5,done5,shot5['serial_no'])

# submit shot6
body={'prompt':shot6['prompt'],'model':MODEL,'aspect_ratio':ASPECT_RATIO,'resolution':RESOLUTION,'quality':QUALITY,'count':1,'input_images':[data_url(make_ref(shot6['ref']))]}
print(json.dumps({'event':'submit','index':6,'id':shot6['id'],'ref':shot6['ref'].name},ensure_ascii=False),flush=True)
resp=request_json('POST','/mcp/api/generate/image',body,180)
serial=serial_from(resp)
if not serial: raise RuntimeError('submit response missing serial_no: '+redact(resp))
print(json.dumps({'event':'submitted','index':6,'serial_no':serial},ensure_ascii=False),flush=True)
done6=poll(serial,900)
save_done(shot6,done6,serial)
print('FINAL_JSON_START')
print(json.dumps({'completed':load_manifest().get('completed'),'failed':load_manifest().get('failed'),'manifest':str(MANIFEST)},ensure_ascii=False,indent=2))
print('FINAL_JSON_END')