#!/usr/bin/env python3
"""三引擎裸描述能力测试: 4 图 × 八维强制框架 × temp 0.1。
输出 results/nl_<arm>.json。盲测匿名映射存 secrets.json（不进报告页）。
"""
import json, base64, time, sys, os, urllib.request
from pathlib import Path

EXP = Path('/home/kai/shared/qwen-eye-exp/nltest')
RES = EXP / 'results'
RES.mkdir(exist_ok=True)
ARM = sys.argv[1]

IMAGES = ['L1_ocr', 'L2_scene', 'L3_night', 'L4_grid']

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

for k in list(os.environ):
    if 'proxy' in k.lower():
        del os.environ[k]

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

# ── 臂实现 ──
QWEN_API = "http://127.0.0.1:8125/v1/chat/completions"

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

def call_qwen(img_path):
    body = json.dumps({"model": "qwen3.8", "temperature": 0.1, "max_tokens": 6000,
        "chat_template_kwargs": {"enable_thinking": False},
        "messages": [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64(img_path)}},
            {"type": "text", "text": PROMPT}]}]}).encode()
    last_err = None
    for attempt in range(4):
        if not _wait_qwen_alive():
            raise RuntimeError(f"qwen-eye dead >600s after {attempt+1} failures")
        req = urllib.request.Request(QWEN_API, data=body, headers={"Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=900) as r:
                resp = json.load(r)
            txt = (resp['choices'][0]['message'].get('content') or '').strip()
            usage = resp.get('usage', {})
            return txt, {"prompt_tokens": usage.get('prompt_tokens'), "completion_tokens": usage.get('completion_tokens')}
        except Exception as e:
            last_err = e
            print(f"  [retry] qwen call failed ({repr(e)[:80]}), waiting for auto-restart...", flush=True)
    raise last_err

ZHIPU_EP = "https://open.bigmodel.cn/api/paas/v4/chat/completions"

def zhipu_key():
    if os.environ.get('ZHIPU_API_KEY'):
        return os.environ['ZHIPU_API_KEY']
    st = json.loads(open('/home/kai/.claude-glm/settings.json').read())
    return st['env']['ANTHROPIC_AUTH_TOKEN']

def call_zhipu(model, img_path):
    body = json.dumps({"model": model, "temperature": 0.1, "max_tokens": 6000,
        "messages": [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64(img_path)}},
            {"type": "text", "text": PROMPT}]}]}).encode()
    req = urllib.request.Request(ZHIPU_EP, data=body, headers={
        "Content-Type": "application/json", "Authorization": f"Bearer {zhipu_key()}"})
    with urllib.request.urlopen(req, timeout=600) as r:
        resp = json.load(r)
    txt = (resp['choices'][0]['message'].get('content') or '').strip()
    usage = resp.get('usage', {})
    return txt, {"prompt_tokens": usage.get('prompt_tokens'), "completion_tokens": usage.get('completion_tokens')}

ARMS = {
    'local':  lambda p: call_qwen(p),
    'hosted': lambda p: call_zhipu('glm-5.3-flash', p),
    '46v':    lambda p: call_zhipu('glm-4.6v', p),
}

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

out = {}
for name in IMAGES:
    t0 = time.time()
    try:
        txt, usage = ARMS[ARM](EXP / f'images/{name}.jpg')
        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}
    except Exception as e:
        out[name] = {'error': repr(e)[:300], 'elapsed_s': round(time.time()-t0, 1)}
    print(f"[{ARM}] {name}: {'ok ' + str(out[name].get('raw_len','?')) + ' chars ' + str(out[name].get('elapsed_s')) + 's' if 'error' not in out[name] else 'ERROR ' + out[name]['error'][:120]}", flush=True)

json.dump(out, open(RES / f'nl_{ARM}.json', 'w'), ensure_ascii=False, indent=1)
print(f"NL_DONE {ARM}")