#!/usr/bin/env python3
"""Bonsai 2 27B 臂: 复用 qwen-eye NL 盲测同题 (4图×八维JSON) + 打工人瞄5帧视频理解。
跑在 http://127.0.0.1:8127 (llama-server, PrismML fork)。
用法: python3 bonsai_arm.py [nl|video|both]
"""
import json, base64, time, re, sys, os, urllib.request
from pathlib import Path

API = "http://127.0.0.1:8127/v1/chat/completions"
EXP = Path('/home/kai/shared/qwen-eye-exp')
OUT = Path('/home/kai/shared/bonsai-vs-qweneye/results')
OUT.mkdir(parents=True, exist_ok=True)

NL_PROMPT = """你是专业的影视画面分析师。请对这张图片做**尽可能详尽**的八维描述（这是AIGC生产管线消费的指标）。严格按以下JSON输出（无markdown围栏）:
{"overall":"整体概览(这张图是什么,一句话概括画面内容)",
"subject":"主体描述(是谁/什么,外观细节:服装/表情/姿态/朝向)",
"layout":"方位布局(主体在画面何处,四象限/三分位;其他元素分布)",
"colors":"颜色(主色及其占比感,配色关系,色调冷暖)",
"lighting":"光照(光源方向与性质,软硬光,阴影形态,明暗对比)",
"composition":"构图(景别估计:远全中近特;构图法则;视觉重心)",
"text":"文字转录(图中任何可读文字原样转录;无则空串)",
"defects":"缺陷检查(畸形/糊化/黑屏/水印/文字乱码等;无则'无')"}
每一维都不要省略字数,写具体。"""

VIDEO_OBS = """以下是同一视频按时间序的5帧。逐帧用自然语言描述你实际看到的:
1. 主体是什么(人/动物/物体), 什么颜色外观?
2. 主体位置(左/中/右)和大小占比?
3. 从第1帧到第5帧, 主体大小/位置/朝向/背景有何变化?
4. 有无明显画面缺陷(畸形/糊化/黑屏/水印)?
5. 综合: 景别(远/全/中/近/特写)? 运镜(固定/跟随/推进/拉远)?
直接描述,不要JSON。"""

VIDEO_SUM_TMPL = """你是影视质检员兼内容分析师。下面是【镜头分析员对视频5帧的独立观察】(自然语言, 唯一事实)。
视频来自B站合集《打工瞄》, 本集标题: 《老板让你帮忙搬东西时的你》(竖屏短视频)。

请输出:
1. **内容梗概**: 这30秒讲了什么? 角色是谁? 情绪/笑点在哪?
2. **运镜分析**: 有哪些镜头切换? 每段运镜方式?
3. **构图分析**: 主体位置/画面平衡/字幕条位置?
4. **一致性**: 角色外观跨帧是否一致? 场景是否连贯?
5. **缺陷**: 畸形/糊化/黑屏/水印/字幕错位?
6. **标题匹配度**: 内容与标题《老板让你帮忙搬东西时的你》是否相符? 简短理由。
严格输出JSON(无markdown):
{{"summary":"内容梗概2-3句","shots":["逐镜头: 00:0X-00:0X 景别/运镜/内容"],"camera":"运镜综述1句","composition":"构图综述1句","consistency":"角色/场景一致性综述","defects":"缺陷或无","title_match":bool,"title_reason":"1句","overall_quality":"1-10分+1句理由"}}

观察: {obs}"""

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]

def call(msgs, n=6000, timeout=900):
    body = json.dumps({"model": "bonsai2", "temperature": 0.1, "max_tokens": n,
        "messages": msgs}).encode()
    req = urllib.request.Request(API, data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        resp = json.load(r)
    txt = (resp['choices'][0]['message'].get('content') or '').strip()
    return txt, resp.get('usage', {})

def parse_json_loose(t):
    m = re.search(r'\{.*\}', t, re.S)
    if not m: return None
    try: return json.loads(m.group(0))
    except Exception: return None

def wait_alive(max_s=300):
    dl = time.time() + max_s
    while time.time() < dl:
        try:
            with urllib.request.urlopen("http://127.0.0.1:8127/health", timeout=5) as r:
                if b'"ok"' in r.read(): return True
        except Exception: pass
        time.sleep(5)
    return False

mode = sys.argv[1] if len(sys.argv) > 1 else "both"
wait_alive()

if mode in ("nl", "both"):
    out = {}
    for name in ['L1_ocr', 'L2_scene', 'L3_night', 'L4_grid']:
        t0 = time.time()
        try:
            txt, usage = call([{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64(EXP / f'nltest/images/{name}.jpg')}},
                {"type": "text", "text": NL_PROMPT}]}])
            parsed = parse_json_loose(txt)
            out[name] = {'raw_len': len(txt), 'elapsed_s': round(time.time()-t0, 1),
                         'usage': usage, 'json_ok': parsed is not None,
                         'dims': parsed, 'raw': txt if parsed is None else None}
            print(f"[nl] {name}: {out[name]['raw_len']} chars, {out[name]['elapsed_s']}s, json={out[name]['json_ok']}", flush=True)
        except Exception as e:
            out[name] = {'error': repr(e)[:300], 'elapsed_s': round(time.time()-t0, 1)}
            print(f"[nl] {name}: ERROR {repr(e)[:150]}", flush=True)
    json.dump(out, open(OUT / 'nl_bonsai.json', 'w'), ensure_ascii=False, indent=1)
    print("saved nl_bonsai.json", flush=True)

if mode in ("video", "both"):
    t0 = time.time()
    frames = sorted((EXP / 'frames_catvideo').glob('frame_*.jpg'))
    msgs = [{"role": "user", "content": []}]
    for f in frames:
        msgs[0]["content"].append({"type": "text", "text": f"【帧 {f.stem.replace('frame_','').replace('s','')}秒】"})
        msgs[0]["content"].append({"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64(f)}})
    msgs[0]["content"].append({"type": "text", "text": VIDEO_OBS})
    obs_txt, u1 = call(msgs, n=3000)
    print(f"[video] observation done: {len(obs_txt)} chars, {round(time.time()-t0,1)}s", flush=True)

    t1 = time.time()
    sum_txt, u2 = call([{"role": "user", "content": VIDEO_SUM_TMPL.format(obs=obs_txt)}], n=3000)
    parsed = parse_json_loose(sum_txt)
    print(f"[video] summary done: json={parsed is not None}, {round(time.time()-t1,1)}s", flush=True)
    json.dump({'observation': obs_txt, 'summary_raw': sum_txt, 'summary': parsed,
               'obs_usage': u1, 'sum_usage': u2, 'total_s': round(time.time()-t0, 1)},
              open(OUT / 'video_bonsai.json', 'w'), ensure_ascii=False, indent=1)
    print("saved video_bonsai.json", flush=True)

print("BONSAI_ARM_COMPLETE")