#!/usr/bin/env python3
"""r054 读图/读视频双模型对比。4道题, 双臂 (Bonsai :8127 / qwen-eye :8125)。
真值锚: manifest_r054_pilot_v3.json h3_prompt + r054_gpr_evidence_v4_camera.json (bg-anchor 仪器级定谳)
"""
import json, base64, time, re, os, urllib.request, sys
from pathlib import Path

FR = Path('/home/kai/shared/bonsai-vs-qweneye/r054_frames')
OUT = Path('/home/kai/shared/bonsai-vs-qweneye/results')
OUT.mkdir(parents=True, exist_ok=True)
GT_PROMPT = open('/home/kai/shared/bonsai-vs-qweneye/r054_gt_prompt.txt').read()

def b64(p): return base64.b64encode(Path(p).read_bytes()).decode()
for k in list(os.environ):
    if 'proxy' in k.lower(): del os.environ[k]

# ── 四道题 ──
Q1 = """【读图题·单帧】这是3D动画的一个镜头帧。请精确回答:
1. 画面里有哪几个角色/生物? 各自在画面什么位置(左/中/右,前景/中景/背景)?
2. 每个角色的外观特征(颜色/身体结构/盔甲或绷带等细节)?
3. 地面是什么? 光线色调如何?
4. 有无文字/水印? 原样转录。
直接描述,不要JSON。"""

Q2 = """【读视频题·8帧时间线】以下是同一镜头按时间顺序的8帧(2.33秒,30fps中的8个采样点)。
这个镜头里发生了什么? 请回答:
1. 有哪些角色? 初始状态如何?
2. 随时间推进,每个角色的动作变化(谁动了?往哪个方向?幅度多大?)?
3. 蜈蚣(红色多足长虫)有没有出现? 它做了什么?
4. 相机有没有运动(推拉/平移/缩放/固定)? 依据是什么(背景/前景如何变化的)?
直接描述,不要JSON。"""

Q3 = """【仪器仲裁题】同一个问题,请只依据画面证据回答,置信度高中低:
这台相机在整段镜头里有没有发生推拉(缩放)或平移(pan)?
如果 affirmed 运动发生了,请给出:发生时刻、方向、幅度估计(px)。
重点观察背景物体(树干/植被)相对画面边缘的位置变化。直接回答。"""

Q4 = """【反写题】如果要把这个镜头原样复现成一个视频生成prompt(英文),请写出:
1. 场景与角色外观描述
2. 按时间顺序的动作节拍(beat 1 / beat 2)
3. 运镜描述
4. 光线与风格
只输出prompt正文,英文。"""

def call(arm, msgs, n, timeout=900):
    if arm == 'bonsai':
        api, model, extra = "http://127.0.0.1:8127/v1/chat/completions", "b", {}
    else:
        api, model, extra = "http://127.0.0.1:8125/v1/chat/completions", "qwen3.8", {"chat_template_kwargs": {"enable_thinking": False}}
    body = json.dumps({"model": model, "temperature": 0.1, "max_tokens": n, "messages": msgs, **extra}).encode()
    req = urllib.request.Request(api, data=body, headers={"Content-Type": "application/json"})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as r:
        resp = json.load(r)
    m = resp['choices'][0]['message']
    txt = (m.get('content') or '').strip()
    if not txt and m.get('reasoning_content'):
        txt = "[EMPTY-CONTENT reasoning_only] " + (m.get('reasoning_content') or '')[:400]
    return txt, round(time.time()-t0,1), resp.get('usage',{})

def frames_msgs(names, prefix):
    msgs = [{"role":"user","content":[]}]
    for nm in names:
        msgs[0]["content"].append({"type":"text","text":f"{prefix}{nm}"})
        msgs[0]["content"].append({"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"+b64(FR/nm)}})
    return msgs

def wait(arm):
    port = 8127 if arm=='bonsai' else 8125
    for _ in range(60):
        try:
            with urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=4) as r:
                if b'ok' in r.read(): return True
        except Exception: pass
        time.sleep(5)
    return False

arm = sys.argv[1]
wait(arm)
names = sorted(os.listdir(FR))
res = {}

# Q1 single frame
t0=time.time()
txt,el,u = call(arm, [{"role":"user","content":[
    {"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"+b64(FR/'f0_50ms.jpg')}},
    {"type":"text","text":Q1}]}], n=6000)
res['q1'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[{arm}] Q1 done {el}s", flush=True)

# Q2 8-frame timeline
msgs = frames_msgs(names, "【帧 ")
msgs[0]["content"].append({"type":"text","text":Q2})
txt,el,u = call(arm, msgs, n=6000)
res['q2'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[{arm}] Q2 done {el}s", flush=True)

# Q3 instrument arbitration
msgs = frames_msgs(names, "【帧 ")
msgs[0]["content"].append({"type":"text","text":Q3})
txt,el,u = call(arm, msgs, n=4000)
res['q3'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[{arm}] Q3 done {el}s", flush=True)

# Q4 reverse prompt writing (no image, from accumulated understanding? NO - fair test: give 8 frames again)
msgs = frames_msgs(names, "【帧 ")
msgs[0]["content"].append({"type":"text","text":Q4})
txt,el,u = call(arm, msgs, n=8000)
res['q4'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[{arm}] Q4 done {el}s", flush=True)

json.dump(res, open(OUT/f'r054_{arm}.json','w'), ensure_ascii=False, indent=1)
print(f"{arm.upper()}_R054_COMPLETE", flush=True)