import re, math, sys
from PIL import Image
from mlx_vlm import load, generate
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import load_config

MODEL = "mlx-community/UI-TARS-1.5-7B-6bit"

def smart_resize(h, w, factor=28, min_pixels=3136, max_pixels=1000000):
    h_bar = max(factor, round(h/factor)*factor)
    w_bar = max(factor, round(w/factor)*factor)
    if h_bar*w_bar > max_pixels:
        beta = math.sqrt((h*w)/max_pixels)
        h_bar = max(factor, math.floor(h/beta/factor)*factor)
        w_bar = max(factor, math.floor(w/beta/factor)*factor)
    elif h_bar*w_bar < min_pixels:
        beta = math.sqrt(min_pixels/(h*w))
        h_bar = math.ceil(h*beta/factor)*factor
        w_bar = math.ceil(w*beta/factor)*factor
    return h_bar, w_bar

orig = Image.open("shot.png").convert("RGB")
ow, oh = orig.size
rh, rw = smart_resize(oh, ow)
orig.resize((rw, rh)).save("shot_resized.png")
print(f"원본 {ow}x{oh} → 리사이즈 {rw}x{rh}", flush=True)

model, processor = load(MODEL)
try:
    config = load_config(MODEL)
except Exception:
    config = model.config

# (label, ground-truth fractional center in original image)
targets = [
    ("the blue '새로 시작' start button", 0.77, 0.58),
    ("the '취소' cancel button", 0.57, 0.58),
    ("the '전송' send button at bottom right", 0.90, 0.94),
    ("the settings gear icon at top right", 0.95, 0.03),
]

for label, gtfx, gtfy in targets:
    q = ("Output only the coordinate of one point in your response. "
         f"What element matches the following task: {label}")
    try:
        prompt = apply_chat_template(processor, config, q, num_images=1)
    except TypeError:
        prompt = apply_chat_template(processor, config, q)
    try:
        out = generate(model, processor, prompt, image=["shot_resized.png"],
                       verbose=False, max_tokens=64)
    except TypeError:
        out = generate(model, processor, prompt, ["shot_resized.png"],
                       verbose=False, max_tokens=64)
    text = out if isinstance(out, str) else getattr(out, "text", str(out))
    m = re.findall(r"(\d+)\s*,\s*(\d+)", text)
    line = f"[{label}] raw={text.strip()[:70]!r}"
    if m:
        x, y = int(m[0][0]), int(m[0][1])
        fx, fy = x/rw, y/rh
        dist = math.hypot((fx-gtfx)*ow, (fy-gtfy)*oh)
        ok = "OK" if dist < 60 else ("~near" if dist < 120 else "MISS")
        line += f" -> ({x},{y}) frac=({fx:.2f},{fy:.2f}) GT=({gtfx:.2f},{gtfy:.2f}) err={dist:.0f}px {ok}"
    else:
        line += " -> NO COORD"
    print(line, flush=True)