#!/usr/bin/env python3
"""Continue GPT Image 2 scene 01 jobs sequentially; no local image editing/processing."""
import base64,json,mimetypes,os,pathlib,re,time,urllib.error,urllib.request
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=(ROOT/'prompts/scene-01-atmosphere-redo-20260522-161016.md').read_text(encoding='utf-8').split('## Prompt',1)[1].split('## 参数',1)[0].strip()
REFS=[ROOT/'outputs/scene-01-base/scene01-ten-coin-staggered-layout-feathered-20260522-135944.png']+sorted((ROOT/'assets/reference-coins/raw').glob('*.jpg'))

def req(method,path,key,body=None,timeout=180):
    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')
    with urllib.request.urlopen(urllib.request.Request(BASE+path,data=data,method=method,headers=headers),timeout=timeout) as r:
        return json.loads(r.read().decode('utf-8',errors='replace'))
def data_obj(x): return x.get('data',x) if isinstance(x,dict) else {}
def serial(x):
    d=data_obj(x); return d.get('serial_no') or d.get('serialNo') or d.get('serial') or x.get('serial_no')
def download_first(sn,name,key):
    deadline=time.time()+1800
    while time.time()<deadline:
        d=data_obj(req('GET',f'/mcp/api/task/{sn}',key,timeout=60)); st=str(d.get('status'))
        print(json.dumps({'event':'poll','name':name,'serial_no':sn,'status':st,'text':d.get('status_text')},ensure_ascii=False),flush=True)
        if st=='2':
            urls=d.get('image_urls') or d.get('images') or d.get('urls') or []
            if isinstance(urls,str): urls=[urls]
            if not urls: return {'name':name,'serial_no':sn,'status':'completed_no_url'}
            with urllib.request.urlopen(urllib.request.Request(urls[0],headers={'User-Agent':'HermesVideo/1.0'}),timeout=180) as r:
                content=r.read()
            dest=OUTDIR/name
            dest.write_bytes(content)
            return {'name':name,'serial_no':sn,'status':'completed','file':str(dest),'bytes':len(content)}
        if st in ('3','4'):
            return {'name':name,'serial_no':sn,'status':'failed','text':d.get('fail_msg') or d.get('status_text')}
        time.sleep(10)
    return {'name':name,'serial_no':sn,'status':'timeout'}
def data_url(p): return 'data:%s;base64,%s'%((mimetypes.guess_type(str(p))[0] or 'image/png'),base64.b64encode(p.read_bytes()).decode('ascii'))
def submit_gpt(name,key):
    imgs=[data_url(p) for p in REFS]
    body={'prompt':PROMPT,'model':'gpt-image-2','aspect_ratio':'3:4','resolution':'1K','quality':'high','count':1,'input_images':imgs}
    sn=serial(req('POST','/mcp/api/generate/image',key,body,timeout=180))
    print(json.dumps({'event':'submitted','name':name,'serial_no':sn},ensure_ascii=False),flush=True)
    return sn

def main():
    key=os.environ.get('MX_AI_API_KEY')
    if not key: raise SystemExit('MX_AI_API_KEY missing')
    results=[]
    # Continue already submitted GPT #2 first, then only after it completes submit GPT #3.
    r2=download_first('2057738216814940160','gpt2image-02.png',key); results.append(r2)
    if r2.get('status')=='completed':
        print(json.dumps({'event':'interval_sleep','seconds':30},ensure_ascii=False),flush=True); time.sleep(30)
        sn3=submit_gpt('gpt2image-03.png',key)
        results.append(download_first(sn3,'gpt2image-03.png',key))
    summary={'ok':True,'outdir':str(OUTDIR),'results':results}
    (OUTDIR/'mxai-generation-summary-gpt-continued.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()