import json, base64, time, 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')

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]

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(msgs, n, timeout=1200):
    api = "http://127.0.0.1:8127/v1/chat/completions"
    body = json.dumps({"model":"b","temperature":0.1,"max_tokens":n,"messages":msgs}).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

# resume: q2/q3/q4 with downscaled frames (720p) to fit ctx
import subprocess
small = Path('/home/kai/shared/bonsai-vs-qweneye/r054_frames_small')
small.mkdir(exist_ok=True)
for f in sorted(FR.glob('*.jpg')):
    subprocess.run(["bash","-c",f"ffmpeg -y -v error -i {f} -vf scale=1280:-2 -q:v 4 {small}/{f.name}"], check=True)

names = sorted(os.listdir(small))
res = {}

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


# Q1 first (lost in crash)
q1 = """【读图题·单帧】这是3D动画的一个镜头帧。请精确回答:
1. 画面里有哪几个角色/生物? 各自在画面什么位置(左/中/右,前景/中景/背景)?
2. 每个角色的外观特征(颜色/身体结构/盔甲或绷带等细节)?
3. 地面是什么? 光线色调如何?
4. 有无文字/水印? 原样转录。
直接描述,不要JSON。"""
txt,el,u = call([{"role":"user","content":[
    {"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"+b64(small/'f0_50ms.jpg')}},
    {"type":"text","text":q1}]}], 6000)
res['q1'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[bonsai] Q1 redo {el}s", flush=True)

# Q2
msgs = small_msgs(); msgs[0]["content"].append({"type":"text","text":Q2})
txt,el,u = call(msgs, 6000)
res['q2'] = {'answer':txt,'elapsed_s':el,'usage':u,'note':'frames downscaled 720p'}
print(f"[bonsai] Q2 retry done {el}s", flush=True)

# Q3
msgs = small_msgs(); msgs[0]["content"].append({"type":"text","text":Q3})
txt,el,u = call(msgs, 4000)
res['q3'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[bonsai] Q3 retry done {el}s", flush=True)

# Q4
msgs = small_msgs(); msgs[0]["content"].append({"type":"text","text":Q4})
txt,el,u = call(msgs, 8000)
res['q4'] = {'answer':txt,'elapsed_s':el,'usage':u}
print(f"[bonsai] Q4 retry done {el}s", flush=True)

json.dump(res, open(OUT/'r054_bonsai.json','w'), ensure_ascii=False, indent=1)
print("BONSAI_R054_COMPLETE", flush=True)