import base64, json, sys, urllib.request
KEY=open('/home/lkraven/.config/litellm/infra-ops-key').read().strip()
SP=sys.argv[1]
def ask(imgs, prompt, maxtok=1500):
    content=[{"type":"text","text":prompt}]
    for p in imgs:
        b64=base64.b64encode(open(f"{SP}/{p}","rb").read()).decode()
        content.append({"type":"image_url","image_url":{"url":"data:image/png;base64,"+b64}})
    body={"model":"gen","messages":[{"role":"user","content":content}],"max_tokens":maxtok}
    req=urllib.request.Request("http://10.250.50.70:4000/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Authorization":"Bearer "+KEY,"Content-Type":"application/json"})
    d=json.loads(urllib.request.urlopen(req,timeout=300).read().decode(),strict=False)
    c=d["choices"][0]
    return (c["message"].get("content") or ""), c.get("finish_reason"), d.get("usage",{}).get("prompt_tokens")

print("===== T4b OCCLUSION (tight prompt, the part that got truncated)")
print("  GROUND TRUTH: a green star is PARTLY HIDDEN BEHIND the grey rectangle, top-right")
o,f,p = ask(["spatial.png"], "One sentence only. Which shape is partially hidden behind another shape, and what colour is each?")
print(f"  [finish={f}] {o.strip()[:400]}")

print("\n===== T4c COLOUR TRAP (the rectangle is NOT square: 160x140)")
print("  GROUND TRUTH: wider than tall")
o,f,p = ask(["spatial.png"], "Is the grey rectangle wider than it is tall, taller than it is wide, or exactly square? Answer in one short sentence.")
print(f"  [finish={f}] {o.strip()[:300]}")

print("\n===== T6 FOUR IMAGES (the seat caps at --limit-mm-per-prompt image:4)")
print("  GROUND TRUTH: 4 images accepted; ocr=text, count=shapes, chart=bar chart, spatial=star+rect+circle")
try:
    o,f,p = ask(["ocr.png","count.png","chart.png","spatial.png"],
                "You are given four images. Name what each one is, in order, one short line each. Nothing else.")
    print(f"  [prompt_tokens={p} finish={f}]")
    for l in o.strip().splitlines()[:8]: print("   ", l)
except Exception as e:
    print("  ERROR:", str(e)[:200])

print("\n===== T7 FIVE IMAGES (should be REJECTED by the seat's cap)")
try:
    o,f,p = ask(["ocr.png","count.png","chart.png","spatial.png","ocr.png"], "How many images did I send?")
    print(f"  accepted (cap not enforced?): {o.strip()[:150]}")
except Exception as e:
    print("  correctly rejected:", str(e)[:160])
