#!/usr/bin/env python3
"""Generate v2 clean coin character setting images.
Correction from user: v1 looked too antique/broken. V2 prioritizes the user's cartoon card mood,
uses original artifact refs only for silhouette and key deep marks, and simplifies rough grooves/patina.
No base64 payloads are persisted to disk.
"""
import base64, json, mimetypes, os, pathlib, re, time, urllib.error, urllib.request
from PIL import Image, ImageDraw

BASE = "https://mcp.mxai.cn"
PROJECT = pathlib.Path("/Users/bot1/Volumes/root_for_ai/AI工作区/通用_产品宣传视频_古钱币杜邦纸钱袋包_20260530_1702")
STYLE_ROOT = PROJECT / "assets/character_style_refs/from_user_sheet_20260605/character_crops"
RAW_ROOT = PROJECT / "assets/reference-coins/raw"
OUTDIR = PROJECT / "outputs/coin-character-clean-set/v2_cartoon_simplified_20260605"
PROMPT_DIR = OUTDIR / "prompts"

ITEMS = [
    {"id":"01_broken_disc", "title":"Broken Disc", "style":"01_broken_disc_character_crop.png", "artifact":"coin-10-cracked-gold-disc.jpg", "shape":"round golden disc coin with one clear deep central crack mark"},
    {"id":"02_ancient_script_oval", "title":"Ancient Script Oval", "style":"02_ancient_script_oval_character_crop.png", "artifact":"coin-09-black-oval-inscribed.jpg", "shape":"dark oval coin with one clear raised ancient-script mark"},
    {"id":"03_scroll_bar", "title":"Scroll Bar", "style":"03_scroll_bar_character_crop.png", "artifact":"coin-06-round-square-hole-dark.jpg", "shape":"round coin with central square opening and a few simplified carved marks"},
    {"id":"04_green_disc", "title":"Green Disc", "style":"04_green_disc_character_crop.png", "artifact":"coin-05-round-square-hole-green.jpg", "shape":"green round coin with central square opening and a few softened raised characters"},
    {"id":"05_central_script_disc", "title":"Central Script Disc", "style":"05_central_script_disc_character_crop.png", "artifact":"coin-07-round-square-hole-ornate.jpg", "shape":"blue-green round coin with central square opening and simplified raised script around it"},
    {"id":"06_bronze_script_hollow", "title":"Bronze Script Hollow", "style":"06_bronze_script_hollow_character_crop.png", "artifact":"coin-08-round-square-hole-brown-ornate.jpg", "shape":"warm bronze round coin with central square opening and simplified raised script"},
    {"id":"07_spanner_bar", "title":"Spanner Bar", "style":"07_spanner_bar_character_crop.png", "artifact":"coin-01-bridge-shaped-gray-metal.jpg", "shape":"gray bridge-shaped ancient money silhouette with rounded toy-like edges"},
    {"id":"08_inscribed_ax", "title":"Inscribed Ax", "style":"08_inscribed_ax_character_crop.png", "artifact":"coin-03-small-spade-dark.jpg", "shape":"small dark spade-shaped ancient money silhouette with one top hole and simplified engraved marks"},
    {"id":"09_textured_bar", "title":"Textured Bar", "style":"09_textured_bar_character_crop.png", "artifact":"coin-04-flat-spade-green-patina.jpg", "shape":"green flat spade-shaped ancient money silhouette with long handle and simplified mineral color patches"},
]

def read_key():
    key = os.environ.get("MX_AI_API_KEY")
    if key: return key.strip()
    for p in [pathlib.Path.home()/".hermes/secret-vault/shared-services.env", pathlib.Path.home()/".hermes/profiles/video/.env"]:
        if p.exists():
            m = re.search(r"^\s*MX_AI_API_KEY\s*=\s*['\"]?([^'\"\n]+)", p.read_text(encoding='utf-8', errors='ignore'), re.M)
            if m: return m.group(1).strip()
    raise RuntimeError("MX_AI_API_KEY missing")

def redact(s,key=''):
    s=str(s)
    if key: s=s.replace(key,'********')
    s=re.sub(r"nb_[A-Za-z0-9_\-]+", "nb_********", s)
    s=re.sub(r"([?&](?:auth|token|sign|signature|x-[^=]+)=)[^&\s]+", r"\1...", s, flags=re.I)
    return s

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

def request_json(method,path,key,body=None,timeout=150):
    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,key)}")
    except Exception as e:
        raise RuntimeError(f"{type(e).__name__} {path}: {redact(e,key)}")

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]
    return None

def urls_from(obj):
    urls=obj.get('image_urls') or obj.get('images') or obj.get('urls') or obj.get('image') or []
    if isinstance(urls,str): urls=[urls]
    return urls

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

def make_prompt(item):
    return f"""
Create a clean cartoon character design image for one anthropomorphic ancient coin character for a warm product video.

Reference use:
- Image 1 is the main style reference. Match its cute premium cartoon feeling, rounded short arms and legs, simplified 3D toy-like rendering, warm clean atmosphere, and friendly proportions.
- Image 2 is only a shape-and-key-mark reference. Use it to preserve the artifact silhouette and the most important deep identity marks only. Do not copy dirty antique roughness, heavy corrosion, many tiny grooves, broken crumbs, dust, or overly realistic old surface.

Character: {item['shape']}.

Render one full-body character centered on a clean warm off-white background. The ancient coin shape is the main torso form. Add only simple short rounded arms and simple short rounded legs, consistent with the user's cartoon character card. The surface should be simplified, smoother, cleaner, and more playful: keep the key dark grooves/cracks/holes/inscriptions as clear graphic marks, but soften and reduce small rough details. It should feel like a cultural-creative toy character made from the coin, not a damaged archaeological relic.

Mood and style: cute but premium, warm product-ad atmosphere, soft studio lighting, rounded edges, subtle contact shadow, clean reference sheet, consistent with a bright playful coin pouch video.

Avoid: text, labels, UI, watermark, card border, pedestal, inset images, realistic dirt, excessive corrosion, many messy scratches, crumbly broken edges, scary mood, face, eyes, mouth, costume, extra accessories.
""".strip()

def make_contact(files, outpath):
    thumbs=[]
    for f in files:
        img=Image.open(f).convert('RGB')
        img.thumbnail((260,300))
        c=Image.new('RGB',(300,360),'white')
        c.paste(img,((300-img.width)//2,15))
        d=ImageDraw.Draw(c)
        d.text((10,330),pathlib.Path(f).stem[:34],fill=(0,0,0))
        thumbs.append(c)
    cols=3; rows=(len(thumbs)+cols-1)//cols
    sheet=Image.new('RGB',(cols*300,rows*360),(245,242,235))
    for i,t in enumerate(thumbs): sheet.paste(t,((i%cols)*300,(i//cols)*360))
    sheet.save(outpath, quality=92)

def main():
    key=read_key()
    OUTDIR.mkdir(parents=True, exist_ok=True); PROMPT_DIR.mkdir(parents=True, exist_ok=True)
    results=[]; pending={}
    for item in ITEMS:
        style=STYLE_ROOT/item['style']; artifact=RAW_ROOT/item['artifact']
        prompt=make_prompt(item)
        prompt_path=PROMPT_DIR/f"{item['id']}_prompt.md"
        prompt_path.write_text(prompt+'\n',encoding='utf-8')
        body={"prompt":prompt,"model":"gpt-image-2","aspect_ratio":"3:4","resolution":"1K","quality":"high","count":1,"input_images":[data_url(style),data_url(artifact)]}
        resp=request_json('POST','/mcp/api/generate/image',key,body,timeout=180)
        serial=serial_from(resp)
        rec={"id":item['id'],"title":item['title'],"serial_no":serial,"status":"submitted" if serial else "submit_failed","prompt_path":str(prompt_path),"style_ref":str(style),"artifact_ref":str(artifact)}
        if not serial: rec['error']=redact(resp,key)
        results.append(rec)
        if serial: pending[serial]=rec
        print(json.dumps({"event":"submitted" if serial else "submit_failed","id":item['id'],"serial_no":serial,"error":rec.get('error')},ensure_ascii=False),flush=True)
        time.sleep(1.5)
    manifest_path=OUTDIR/'manifest.json'
    deadline=time.time()+900
    while pending and time.time()<deadline:
        for serial in list(pending):
            rec=pending[serial]
            try:
                obj=nested(request_json('GET',f'/mcp/api/task/{serial}',key,timeout=90))
                st=str(obj.get('status')); txt=obj.get('status_text') or obj.get('message') or obj.get('fail_msg')
                print(json.dumps({"event":"poll","id":rec['id'],"serial_no":serial,"status":st,"text":txt},ensure_ascii=False),flush=True)
                if st=='2':
                    files=[]
                    for i,url in enumerate(urls_from(obj),1):
                        dest=OUTDIR/f"{rec['id']}_cartoon_clean_character_v2_{i}.png"
                        download(url,dest); files.append(str(dest))
                    rec.update({"status":"completed","files":files}); pending.pop(serial,None)
                elif st in {'3','4'}:
                    rec.update({"status":"failed","fail_msg":redact(obj.get('fail_msg') or obj,key)}); pending.pop(serial,None)
            except Exception as e:
                rec['last_error']=redact(e,key)
        completed=[f for r in results if r.get('status')=='completed' for f in r.get('files',[])]
        manifest_path.write_text(json.dumps({"results":results,"completed":len(completed)},ensure_ascii=False,indent=2),encoding='utf-8')
        if pending: time.sleep(5)
    for rec in pending.values(): rec['status']='timeout'
    completed=[f for r in results if r.get('status')=='completed' for f in r.get('files',[])]
    contact=None
    if completed:
        contact=OUTDIR/'cartoon_clean_character_v2_contact_sheet.jpg'
        make_contact(completed,contact)
    manifest_path.write_text(json.dumps({"results":results,"completed":len(completed),"contact_sheet":str(contact) if contact else None},ensure_ascii=False,indent=2),encoding='utf-8')
    print('FINAL_JSON_START')
    print(json.dumps({"ok":True,"outdir":str(OUTDIR),"completed":len(completed),"contact_sheet":str(contact) if contact else None,"manifest":str(manifest_path),"results":results},ensure_ascii=False,indent=2))
    print('FINAL_JSON_END')

if __name__=='__main__': main()