#!/usr/bin/env python3
import argparse, base64, json, mimetypes, os, re, sys, urllib.request, urllib.error
from pathlib import Path
DEFAULT_API_KEY='sk-9ffd188da865aa49-mtv9v6-54d294f2'
def refval(x):
    if x.startswith('http://') or x.startswith('https://') or x.startswith('data:image/'): return x
    p=Path(x); mime=mimetypes.guess_type(str(p))[0] or 'image/png'
    return 'data:%s;base64,%s'%(mime,base64.b64encode(p.read_bytes()).decode())
def save_payload(text,out):
    payloads=[]
    for line in text.splitlines():
        line=line.strip()
        if line.startswith('data:'): line=line[5:].strip()
        if not line or line=='[DONE]': continue
        try: payloads.append(json.loads(line))
        except: pass
    if not payloads:
        payloads=[json.loads(text)]
    for payload in payloads:
        data=payload.get('data') if isinstance(payload,dict) else None
        if not data: continue
        item=data[0] if isinstance(data,list) else data
        if not isinstance(item,dict): continue
        b64=item.get('b64_json') or item.get('base64') or item.get('image_base64')
        url=item.get('url') or item.get('image_url')
        if b64:
            if b64.startswith('data:'): b64=b64.split(',',1)[1]
            Path(out).write_bytes(base64.b64decode(b64)); print(out); return
        if url:
            urllib.request.urlretrieve(url,out); print(out); return
    raise SystemExit('No image in response: '+text[:1000])
def main():
    ap=argparse.ArgumentParser(); ap.add_argument('--prompt',required=True); ap.add_argument('--output','-o',required=True)
    ap.add_argument('--input-ref', action='append', default=[]); ap.add_argument('--image', '-i', action='append', default=[]); ap.add_argument('--image-url', action='append', default=[])
    ap.add_argument('--base-url', default=os.environ.get('GEN_IMG_BASE_URL','http://192.168.40.11:20128')); ap.add_argument('--api-key', default=os.environ.get('GEN_IMG_API_KEY') or DEFAULT_API_KEY)
    ap.add_argument('--model', default='cx/gpt-5.5-image'); ap.add_argument('--n', type=int, default=1); ap.add_argument('--size', default='auto'); ap.add_argument('--quality', default='auto'); ap.add_argument('--background', default='auto'); ap.add_argument('--image-detail', default='high'); ap.add_argument('--output-format', default='png')
    a=ap.parse_args(); refs=[refval(r) for r in (a.input_ref+a.image+a.image_url)]
    body={'model':a.model,'prompt':a.prompt,'n':a.n,'size':a.size,'quality':a.quality,'background':a.background,'image_detail':a.image_detail,'output_format':a.output_format}
    if refs: body['image']=refs[0] if len(refs)==1 else refs
    req=urllib.request.Request(a.base_url.rstrip()+'/v1/images/generations', data=json.dumps(body).encode(), headers={'Content-Type':'application/json','Authorization':'Bearer '+a.api_key,'Accept':'text/event-stream'}, method='POST')
    try:
        with urllib.request.urlopen(req, timeout=300) as r: text=r.read().decode('utf-8','ignore')
    except urllib.error.HTTPError as e:
        print(e.read().decode('utf-8','ignore')[:2000], file=sys.stderr); raise
    save_payload(text,a.output)
if __name__=='__main__': main()