#!/usr/bin/env python3
"""identity 维度重判: 用宴会服 turnaround (银白鱼尾裙) 替换日常装锚点, 重评所有含人物的 fail 帧"""
import json, base64, time, re, urllib.request
from pathlib import Path

OSS = Path('/data/workspace/kais-aigc-platform/data/oss/1785508691757')
IFRAME = OSS / 'p11/iframes'
TR = OSS / 'p04/turnaround_sheets'
API = "http://127.0.0.1:8125/v1/chat/completions"
OUT = Path('/home/kai/shared/shenyuan-iframe-eval/results_identity_recheck.jsonl')

def b64(p): return base64.b64encode(Path(p).read_bytes()).decode()

def call(msgs, n=900, timeout=300):
    body = json.dumps({"model": "qwen3.8", "temperature": 0.1, "max_tokens": n,
        "chat_template_kwargs": {"enable_thinking": False}, "messages": msgs}).encode()
    req = urllib.request.Request(API, data=body, headers={"Content-Type": "application/json"})
    last=None
    for a in range(3):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                resp = json.load(r)
            txt = (resp['choices'][0]['message'].get('content') or '').strip()
            if txt: return txt
        except Exception as e:
            last=e; time.sleep(3)
    raise RuntimeError(f"LLM fail: {last}")

OBS = """图1是宴会厅场景短剧的【条件帧候选】(沈知意=22岁东亚女性,黑长直发,杏仁眼,宴会场景着装=银白色鱼尾礼服裙+银色高跟鞋+珍珠耳环)。
图2是该角色的【宴会服turnaround设定图】。
仔细看两张图, 只回答角色身份一致性:
1. 图1中是否出现人物/人物部位(脸/手/身体)? 是什么?
2. 若可见: 发型发色(黑长直?)、年龄气质、服装(银白鱼尾礼服?)是否与图2设定一致?
3. 若不可见(纯手部剪影/纯环境/纯物品): 回答"人物不可见, 无法比对"。
直接描述, 不超过150字, 不要JSON。"""

JUDGE = """你是质检判定官。【观察】是质检员对条件帧的人物身份一致性描述; 【设定】=沈知意 宴会服: 22岁东亚女性, 黑长直发, 杏仁眼, 银白色鱼尾礼服裙, 银色高跟鞋, 珍珠耳环。
判定 identity_match:
- 观察说"人物不可见" → identity_match=true, verdict原样 (身份维度不适用, 不因不可见而fail)
- 可见且符合(发型/服装/年龄气质基本一致, 服装细节可有小差异) → identity_match=true
- 可见且明确不符(发型/发色/年龄性别错误, 服装完全不同如休闲装vs礼服) → identity_match=false
输出JSON: {"identity_match":bool,"identity_note":"一句话"}

观察: <OBS>"""

def main():
    rows=[json.loads(l) for l in open('/home/kai/shared/shenyuan-iframe-eval/results.jsonl') if l.strip()]
    by_id={}
    for r in rows:
        if 'verdict' in r: by_id[r['id']]=r
    targets=[r for r in by_id.values() if r['verdict'].get('verdict')=='fail' and r['verdict'].get('identity_match') is False]
    print(f"recheck targets: {len(targets)}")
    # banquet TR: tier1 优先 (最新定稿), fallback turnaround_shenzhiyi_banquet
    for cand in ['tier1_shenzhiyi_banquet_v2.png','turnaround_shenzhiyi_banquet.png']:
        tr = TR / cand
        if tr.exists(): print("anchor:", cand); break
    trb64 = b64(tr)
    done=set()
    if OUT.exists():
        for l in OUT.read_text().splitlines():
            try: done.add(json.loads(l)['id'])
            except: pass
    for i,r in enumerate(sorted(targets,key=lambda x:x['id'])):
        if r['id'] in done: continue
        fp = IFRAME / (r['id']+'.png')
        t0=time.time()
        try:
            msgs=[
              {"role":"user","content":[{"type":"image_url","image_url":{"url":"data:image/png;base64,"+b64(fp)}},{"type":"text","text":"图1: 条件帧候选"}]},
              {"role":"user","content":[{"type":"image_url","image_url":{"url":"data:image/png;base64,"+trb64}},{"type":"text","text":"图2: 沈知意宴会服设定图"}]},
              {"role":"user","content":OBS}]
            obs=call(msgs, n=500)
            time.sleep(1)
            j=call([{"role":"user","content":JUDGE.replace('<OBS>',obs)}],n=400)
            mj=re.search(r'\{.*\}',j,re.S)
            v=json.loads(mj.group(0)) if mj else {"identity_match":None,"identity_note":"PARSE_FAIL:"+j[:150]}
            rec={"id":r['id'],"obs":obs,"recheck":v,"elapsed_s":round(time.time()-t0,1)}
        except Exception as e:
            rec={"id":r['id'],"error":str(e)[:200]}
        with open(OUT,'a') as f: f.write(json.dumps(rec,ensure_ascii=False)+'\n')
        print(f"[{i+1}/{len(targets)}] {r['id']}: identity_match={rec.get('recheck',{}).get('identity_match')} {round(time.time()-t0,1)}s",flush=True)
    print("DONE")

if __name__=='__main__':
    main()