import base64, json, sys, urllib.request
KEY=open('/home/lkraven/.config/litellm/infra-ops-key').read().strip()
SP=sys.argv[1]; MODEL=sys.argv[2] if len(sys.argv)>2 else "gen"

def ask(imgs, prompt, maxtok=900):
    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":MODEL,"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")

TESTS=[
 ("T1 OCR (exact strings)", ["ocr.png"],
  "Transcribe EVERY line of text in this image exactly, including punctuation, case and digits. One line per line. Nothing else.",
  "HELIOTROPE-49 / batch 7734 / rev 2b / Expires: 2027-03-14 / lot#  aQ8-zX2-004 / tiny print: verify seal"),
 ("T2 counting + attribute binding", ["count.png"],
  "Count each kind of shape. Answer in exactly three lines: 'yellow circles: N', 'purple triangles: N', 'red squares: N'.",
  "yellow circles: 7, purple triangles: 3, red squares: 1"),
 ("T3 chart reading", ["chart.png"],
  "Read this bar chart. Give the value for every day, then state which day is highest and which is lowest.",
  "Mon 34, Tue 71, Wed 22, Thu 58, Fri 93, Sat 47; highest Fri, lowest Wed"),
 ("T4 spatial + occlusion", ["spatial.png"],
  "Describe the image: every shape, its colour, its position, and state which shape is partially hidden behind another.",
  "green star top-right PARTLY BEHIND a grey rectangle; orange circle lower-left; text 'left side text' upper-left"),
 ("T5 multi-image comparison (2 images)", ["count.png","chart.png"],
  "You are given two images. In one sentence each, say what image 1 shows and what image 2 shows, then state which one contains text.",
  "img1 = scattered shapes (no text); img2 = bar chart (has text)"),
]

for name, imgs, prompt, truth in TESTS:
    try:
        out, fin, ptok = ask(imgs, prompt)
    except Exception as e:
        print(f"\n===== {name}\n  ERROR: {e}"); continue
    leak = "<think>" in out
    print(f"\n===== {name}   [prompt_tokens={ptok} finish={fin} think_leak={leak}]")
    print(f"  GROUND TRUTH: {truth}")
    print("  MODEL:")
    for line in out.strip().splitlines()[:16]:
        print("   ", line)