#!/usr/bin/env python3
"""最终版本地臂: Qwen3.8-27B CPU (:8126)
- 每图一条 user 消息 (绕过 llama.cpp 单消息多图丢失 bug: 5图只算3张)
- thinking off (速度), 自然语言观察 + 文本判定
- 28 任务: 20 正配对 + 8 负对照
"""
import json, base64, time, re, urllib.request
from pathlib import Path

EP = Path("/home/kai/runs/ep-cat-worker-001")
FR = Path("/home/kai/shared/qwen-eye-exp/frames")
OUT = Path("/home/kai/shared/qwen-eye-exp/results")
API = "http://127.0.0.1:8125/v1/chat/completions"
shots = json.load(open(EP / '.pipeline-assets/shot-list.json'))['value']
spec_by_id = {s['shot_id']: s for s in shots}

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

JUDGE_TMPL = """你是影视质检员。下面是【镜头分析员对视频帧的独立观察】(自然语言)和【SPEC规格】。
逐维对比判断, 观察是唯一事实。基准:
- 远景/WS=主体很小环境为主; 全景/FS=全身; 中景/MS=膝上; 近景/MCU=胸上; 过肩/OTS=有前景肩膀; 特写/CU/BCU=头或局部; INSERT=插入物件特写
- 运镜按观察的帧间变化: 主体渐大≈推进; 主体渐小≈拉远; 主体稳+背景动≈跟随; 无变化≈固定
- 三分法=主体在两侧1/3区域(非正中); 引导线=画面有指向主体的线条
- 容差: 景别相邻一级算OK; 观察到猫/动物=SPEC中拟人角色匹配成功(本项目角色是拟人橘猫)
- 缺陷(畸形/糊化/黑屏) → 直接fail
- 多段景别如 MCU→WS 按观察的首尾变化方向判

观察: <OBS>

SPEC: <SPEC>

严格输出JSON(无markdown): {"framing_ok":bool,"camera_ok":bool,"composition_ok":bool,"axis_ok":bool,"content_match":bool,"confidence":0.0-1.0,"verdict":"pass|fail","reasons":"一句话"}"""

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

def call(msgs, n=2000):
    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"})
    for a in range(2):
        try:
            with urllib.request.urlopen(req, timeout=3600) as r:
                resp = json.load(r)
            txt = (resp['choices'][0]['message'].get('content') or '').strip()
            if txt: return txt
        except Exception as e:
            if a == 1: raise
            time.sleep(5)
    return ''

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

def judge(sid, spec):
    msgs = []
    for i, f in enumerate(['f0.jpg', 'f1.jpg', 'f2.jpg', 'f3.jpg', 'f4.jpg']):
        msgs.append({"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64(FR / sid / f)}},
            {"type": "text", "text": f"(帧{i+1}/5)"}]})
    msgs.append({"role": "user", "content": OBS})
    t0 = time.time()
    obs = call(msgs)
    spec_slim = json.dumps({k: spec.get(k) for k in
        ['shot_id', 'shot_type', 'framing', 'camera_movement', 'composition', 'axis_line', 'intent']},
        ensure_ascii=False)
    jtxt = call([{"role": "user", "content": JUDGE_TMPL.replace('<OBS>', obs).replace('<SPEC>', spec_slim)}], n=800)
    c = pj(jtxt)
    return {"video": sid, "spec_of": spec['shot_id'], "elapsed_s": round(time.time() - t0, 1), "obs": obs,
            "verdict": c.get('verdict') if c else "PARSE_FAIL",
            "dims": {k: c.get(k) for k in ['framing_ok', 'camera_ok', 'composition_ok', 'axis_ok', 'content_match']} if c else None,
            "reasons": c.get('reasons') if c else jtxt[:200]}

CONTROLS = [("S1-01", "S4-01"), ("S5-03", "S8-01"), ("S7-03", "S6-02"), ("S4-02", "S2-02"),
            ("S8-01", "S1-03"), ("S3-01", "S1-02"), ("S1-02", "S2-01"), ("S7-02", "S5-03")]

def main():
    log = open(OUT / 'judge_final_local.jsonl', 'a')
    done = set()
    p = OUT / 'judge_final_local.jsonl'
    if p.exists():
        for line in open(p):
            try:
                d = json.loads(line)
                if 'error' not in d:
                    done.add((d.get('video'), d.get('spec_of')))
            except Exception:
                pass
    tasks = [("pos", s['shot_id'], s) for s in shots]
    tasks += [("neg", v, spec_by_id[sp]) for v, sp in CONTROLS]
    tasks = [t for t in tasks if (t[1], t[2]['shot_id']) not in done]
    print(f"resume: {len(done)} done, {len(tasks)} remaining", flush=True)
    for i, (tag, vid, spec) in enumerate(tasks):
        try:
            r = judge(vid, spec); r['tag'] = tag
        except Exception as e:
            r = {"tag": tag, "video": vid, "error": str(e)[:300]}
        log.write(json.dumps(r, ensure_ascii=False) + "\n"); log.flush()
        print(f"[{i+1}/{len(tasks)}] {tag} v={vid} s={r.get('spec_of')} -> {r.get('verdict','ERR')} | {str(r.get('reasons',''))[:70]}", flush=True)
    print("DONE")

if __name__ == '__main__':
    main()