#!/usr/bin/env python3
"""Generate clean anthropomorphic coin character setting images.
Each task uses exactly 2 input refs:
1) cropped user character card/style crop
2) original coin artifact image
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/from_user_sheet_original_artifacts_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", "body":"a cracked golden round disc coin body, the central vertical crack remains clear and recognizable"},
    {"id":"02_ancient_script_oval", "title":"Ancient Script Oval", "style":"02_ancient_script_oval_character_crop.png", "artifact":"coin-09-black-oval-inscribed.jpg", "body":"a dark black-brown oval inscribed coin body, preserving the oval silhouette and raised ancient script symbol"},
    {"id":"03_scroll_bar", "title":"Scroll Bar", "style":"03_scroll_bar_character_crop.png", "artifact":"coin-06-round-square-hole-dark.jpg", "body":"a dark round square-hole coin body with carved symbols and a central square opening, preserving the original dark patina"},
    {"id":"04_green_disc", "title":"Green Disc", "style":"04_green_disc_character_crop.png", "artifact":"coin-05-round-square-hole-green.jpg", "body":"a green patinated round square-hole coin body, preserving its green-gray patina, central square hole, and 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", "body":"a blue-green bronze round square-hole coin body with ornate raised script around the square hole"},
    {"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", "body":"a brown bronze round square-hole coin body with ornate raised script and warm aged metal highlights"},
    {"id":"07_spanner_bar", "title":"Spanner Bar", "style":"07_spanner_bar_character_crop.png", "artifact":"coin-01-bridge-shaped-gray-metal.jpg", "body":"a gray bridge-shaped ancient money body, preserving the curved bar silhouette and metallic surface"},
    {"id":"08_inscribed_ax", "title":"Inscribed Ax", "style":"08_inscribed_ax_character_crop.png", "artifact":"coin-03-small-spade-dark.jpg", "body":"a dark small spade-shaped ancient money body, preserving the top hole, forked lower legs of the artifact shape, and engraved characters"},
    {"id":"09_textured_bar", "title":"Textured Bar", "style":"09_textured_bar_character_crop.png", "artifact":"coin-04-flat-spade-green-patina.jpg", "body":"a green patinated flat spade-shaped ancient money body, preserving its long handle, rectangular lower body, mottled mineral texture and copper-green oxidation"},
]

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():
            text=p.read_text(encoding='utf-8', errors='ignore')
            m=re.search(r"^\s*MX_AI_API_KEY\s*=\s*['\"]?([^'\"\n]+)", text, 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 character design image for one anthropomorphic ancient coin artifact character.

Input reference use:
- Image 1 is only a style reference for the short rounded arms and legs, soft toy-like 3D material rendering, and friendly standing character proportions. Ignore all card text, borders, icons, inset pictures, labels, and pedestal from image 1.
- Image 2 is the strict artifact reference for the body. Preserve the artifact body shape, color, patina, surface texture, cracks, holes, inscriptions, raised relief and proportions from image 2.

Character to render: {item['body']}.

Final image:
One full-body character centered on a clean warm off-white background. The artifact itself is the torso/body. Add only simple short rounded arms and simple short rounded legs in a consistent premium toy-like 3D style, matching the style of image 1. The character stands in a neutral slightly lively pose suitable for later video animation. No face, no eyes, no mouth, no costume, no accessories, no pedestal, no card frame, no text, no labels, no UI, no watermark. Keep the silhouette readable and the original artifact identity obvious.

Style: premium cultural-creative product character design, realistic metal/stone surface, soft studio lighting, subtle contact shadow, high detail, clean reference sheet.
""".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']
        if not style.exists() or not artifact.exists(): raise FileNotFoundError(f"missing refs for {item['id']}: {style} {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)
        if not serial:
            rec={"id":item['id'],"status":"submit_failed","error":redact(resp,key),"prompt_path":str(prompt_path),"style_ref":str(style),"artifact_ref":str(artifact)}
            results.append(rec); print(json.dumps({"event":"submit_failed",**rec},ensure_ascii=False),flush=True); continue
        rec={"id":item['id'],"title":item['title'],"serial_no":serial,"status":"submitted","prompt_path":str(prompt_path),"style_ref":str(style),"artifact_ref":str(artifact)}
        results.append(rec); pending[serial]=rec
        print(json.dumps({"event":"submitted","id":item['id'],"serial_no":serial},ensure_ascii=False),flush=True)
        time.sleep(1.5)
    manifest_path=OUTDIR/'manifest.json'
    manifest_path.write_text(json.dumps({"results":results},ensure_ascii=False,indent=2),encoding='utf-8')
    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))
                status=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":status,"text":txt},ensure_ascii=False),flush=True)
                if status=='2':
                    files=[]
                    for i,url in enumerate(urls_from(obj),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"{rec['id']}_clean_character_{i}{ext}"
                        download(url,dest); files.append(str(dest))
                    rec.update({"status":"completed","files":files})
                    pending.pop(serial,None)
                elif status 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)
        manifest_path.write_text(json.dumps({"results":results},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',[])]
    if completed:
        make_contact(completed, OUTDIR/'clean_character_contact_sheet.jpg')
    manifest_path.write_text(json.dumps({"results":results,"contact_sheet":str(OUTDIR/'clean_character_contact_sheet.jpg') if completed else None},ensure_ascii=False,indent=2),encoding='utf-8')
    print('FINAL_JSON_START')
    print(json.dumps({"ok":True,"outdir":str(OUTDIR),"manifest":str(manifest_path),"completed":len(completed),"contact_sheet":str(OUTDIR/'clean_character_contact_sheet.jpg') if completed else None,"results":results},ensure_ascii=False,indent=2))
    print('FINAL_JSON_END')

if __name__=='__main__':
    main()
