#!/usr/bin/env python3
"""
Eval ScreenSpot-Pro with the privileged-teacher prompt (same as 16c demo).
Batched per BATCH_SIZE samples; writes a log line per batch.
"""
import argparse
import json
import sys
import time
from pathlib import Path

import torch
torch.backends.cuda.enable_cudnn_sdp(False)  # avoid cuDNN frontend errors on some nodes
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
from transformers import AutoProcessor, AutoModelForImageTextToText

from windom.core import (
    GROUNDING_SYSTEM_PROMPT_QWEN,
    parse_tool_call,
    point_in_bbox,
)

DEMO_PREFIX = "\nThis is an example for a response to the question:\n"


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--subfolder", default="",
                    help="HF repo subfolder containing the model files (e.g. 'final')")
    ap.add_argument("--samples-json",
                    default=str(Path(__file__).resolve().parents[4]
                                / "eval_benchmarks" / "screenspot_pro" / "samples.json"))
    ap.add_argument("--n", type=int, default=None)
    ap.add_argument("--max-new-tokens", type=int, default=1024)
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--suffix", default="Now answer with the tool call only:")
    ap.add_argument("--no-demo", action="store_true",
                    help="Mode A: skip the GT-center demo entirely. Use the same "
                         "student-style no-info prompt as own-test inference. "
                         "This is a fair head-to-head with the public SS-Pro "
                         "leaderboard.")
    ap.add_argument("--official-prompt", action="store_true",
                    help="Use the exact prompt from the official ScreenSpot-Pro "
                         "evaluator (likaixin2000/screenspot-pro-gui-grounding, "
                         "models/qwen3_5.py). System role + structured <tools> "
                         "JSON schema + 'guess if unsure' fallback. Implies --no-demo.")
    ap.add_argument("--max-pixels", type=int, default=8294400,
                    help="Max image pixels after smart_resize (when --official-prompt). "
                         "Default 8,294,400 (= 2880x2880, what the official Qwen3.5 "
                         "ScreenSpot-Pro eval uses). Set lower if running OOM.")
    ap.add_argument("--rescale-pixel-to-norm", action="store_true",
                    help="When the model emits PIXEL coords (e.g. argus-official), "
                         "rescale predictions to 0-1000 norm before scoring against "
                         "the GT bbox. Skip this for un-fine-tuned base models or "
                         "models trained with 0-1000 normalized targets — they "
                         "already emit in 0-1000 space.")
    ap.add_argument("--batch-size-log", type=int, default=50,
                    help="Print a batch summary every N samples.")
    ap.add_argument("--out", type=Path, default=None,
                    help="Per-sample JSONL output. Default: <samples-json-dir>/preds.4B-teacher.jsonl")
    args = ap.parse_args(argv)

    samples_json = Path(args.samples_json)
    if args.out is None:
        ckpt_tag = str(args.checkpoint).rstrip("/").split("/")[-1].replace("/", "_")
        args.out = samples_json.parent / f"preds.{ckpt_tag}.teacher.jsonl"

    print(f"Loading: {args.checkpoint}" + (f" (subfolder={args.subfolder})" if args.subfolder else ""), flush=True)
    sub_kw = {"subfolder": args.subfolder} if args.subfolder else {}
    processor = AutoProcessor.from_pretrained(args.checkpoint, **sub_kw)
    model = AutoModelForImageTextToText.from_pretrained(
        args.checkpoint, torch_dtype=torch.bfloat16, device_map=args.device, **sub_kw,
    )
    model.eval()

    data = json.load(open(samples_json))
    if args.n: data = data[:args.n]
    print(f"  loaded {len(data)} samples from {samples_json}", flush=True)
    print(f"  per-sample output: {args.out}", flush=True)
    n_correct = 0; n_no_pred = 0

    # Resume: skip samples already written to --out
    done_idxs = set()
    if Path(str(args.out)).exists():
        for line in open(str(args.out)):
            try:
                r = json.loads(line)
                if r.get("idx") is not None:
                    done_idxs.add(r["idx"])
                    if r.get("click_in_gt"): n_correct += 1
                    if r.get("no_pred"): n_no_pred += 1
            except: pass
        if done_idxs:
            print(f"Resuming: {len(done_idxs)} already done, {len(data)-len(done_idxs)} remaining", flush=True)

    out_f = open(args.out, "a" if done_idxs else "w")

    n_processed = len(done_idxs)
    batch_start_t = time.time()
    batch_correct = 0; batch_np = 0
    for i, s in enumerate(data):
        if i in done_idxs:
            continue
        # ScreenSpot-Pro normalized format:
        #   image_path: "eval_benchmarks/screenspot_pro/images/XXXX.png"
        #   instruction: "Click ..."
        #   gt_bbox: [y1, x1, y2, x2] in 0-1000 (Gemma convention)
        #   gt_point: [y, x] in 0-1000
        img_path = s["image_path"]
        if not Path(img_path).is_absolute():
            img_path = str(Path(__file__).resolve().parent.parent / img_path)
        if not Path(img_path).exists():
            n_no_pred += 1; batch_np += 1
            continue
        try:
            image = Image.open(img_path).convert("RGB")
        except (OSError, Image.UnidentifiedImageError) as e:
            print(f"  [skip idx={i}] image load failed: {e}", flush=True)
            n_no_pred += 1; batch_np += 1
            continue
        instruction = s["instruction"]
        gb = s.get("gt_bbox") or [-1, -1, -1, -1]
        # un-swap [y1,x1,y2,x2] -> [x1,y1,x2,y2]
        gt_xyxy = [gb[1], gb[0], gb[3], gb[2]]
        # demo: tool-call format with center-point click
        cx = int((gt_xyxy[0] + gt_xyxy[2]) / 2.0)
        cy = int((gt_xyxy[1] + gt_xyxy[3]) / 2.0)
        demo = (
            '<tool_call>\n'
            f'{{"name": "computer_use", "arguments": {{"action": "left_click", "coordinate": [{cx}, {cy}]}}}}\n'
            '</tool_call>'
        )

        if args.official_prompt:
            # Use original image dimensions for the system-prompt's resolution
            # claim. We do NOT pre-resize the image — the processor handles
            # its own downsampling/patching, so adding our own smart_resize
            # was overengineering and broke for some torchvision-less envs.
            rw, rh = image.width, image.height
            screen_w, screen_h = rw, rh

            # Exact official ScreenSpot-Pro prompt with the resized W×H injected.
            sys_text = (
                "You are a helpful assistant. The user will give you an instruction, "
                "and you MUST left click on the corresponding UI element via tool call. "
                "If you are not sure about where to click, guess a most likely one.\n\n"
                "# Tools\n\n"
                "You may call one or more functions to assist with the user query.\n\n"
                "You are provided with function signatures within <tools></tools> XML tags:\n"
                "<tools>\n"
                '{"type": "function", "function": {"name": "computer_use", '
                f'"description": "Use a mouse to interact with a computer.\\n* The screen\'s '
                f'resolution is {screen_w}x{screen_h}.\\n* Make sure to click any buttons, '
                'links, icons, etc with the cursor tip in the center of the element. \\n* '
                'You can only use the left_click action to interact with the computer.", '
                '"parameters": {"properties": {"action": {"description": '
                '"The action to perform. The available actions are:\\n* `left_click`: Click '
                'the left mouse button with coordinate (x, y).", '
                '"enum": ["left_click"], "type": "string"}, '
                '"coordinate": {"description": "(x, y): The x (pixels from the left edge) '
                'and y (pixels from the top edge) coordinates to move the mouse to. '
                'Required only by `action=left_click`.", "type": "array"}, '
                '"required": ["action"], "type": "object"}}}\n'
                "</tools>\n\n"
                "For each function call, return a json object with function name and "
                "arguments within <tool_call></tool_call> XML tags:\n"
                "<tool_call>\n"
                '{"name": <function-name>, "arguments": <args-json-object>}\n'
                "</tool_call>"
            )
            messages = [
                {"role": "system",
                 "content": [{"type": "text", "text": sys_text}]},
                {"role": "user",
                 "content": [{"type": "image"},
                             {"type": "text", "text": instruction}]},
            ]
            # Note: under official prompt, predicted coordinates are in
            # PIXEL space of the resized image (0..rw, 0..rh), NOT the
            # 0-1000 normalized space the GT bbox uses. We must rescale
            # predictions back to 0-1000 before scoring against gt_xyxy.
            args._official_pixel_w = rw
            args._official_pixel_h = rh
        elif args.no_demo:
            # Mode A — our own-test prompt format, no GT leak.
            user_text = (f"{GROUNDING_SYSTEM_PROMPT_QWEN}\n\n{instruction}\n"
                         f"Output ONLY the tool call. Do NOT reason, "
                         f"explain, or add anything else.")
            messages = [{
                "role": "user",
                "content": [{"type": "image"}, {"type": "text", "text": user_text}],
            }]
        else:
            user_text = (f"{GROUNDING_SYSTEM_PROMPT_QWEN}\n\n{instruction}"
                         f"{DEMO_PREFIX}{demo}\n{args.suffix}")
            messages = [{
                "role": "user",
                "content": [{"type": "image"}, {"type": "text", "text": user_text}],
            }]
        prompt = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
            enable_thinking=False)
        inputs = processor(text=[prompt], images=[image], return_tensors="pt",
                            padding=False).to(args.device)
        with torch.no_grad():
            gen = model.generate(
                **inputs, max_new_tokens=args.max_new_tokens,
                do_sample=False, eos_token_id=processor.tokenizer.eos_token_id,
            )
        new_tokens = gen[0, inputs["input_ids"].shape[1]:]
        response = processor.tokenizer.decode(new_tokens, skip_special_tokens=False)
        pred = parse_tool_call(response)
        # Optional pixel→norm rescaling: only models explicitly trained to
        # emit pixel coords (argus-official) need this. Un-fine-tuned base
        # / Argus / SFT all emit 0-1000 normalized regardless of the
        # system prompt's resolution claim.
        if args.rescale_pixel_to_norm and pred is not None:
            rw = getattr(args, "_official_pixel_w", 1000)
            rh = getattr(args, "_official_pixel_h", 1000)
            pred = [pred[0] * 1000.0 / rw, pred[1] * 1000.0 / rh]
        ok = point_in_bbox(pred, gt_xyxy)
        if ok: n_correct += 1; batch_correct += 1
        if pred is None: n_no_pred += 1; batch_np += 1
        n_processed += 1

        out_f.write(json.dumps({
            "idx": i,
            "image_path": s["image_path"],
            "instruction": instruction,
            "gt_bbox_xyxy": gt_xyxy,
            "platform": s.get("platform"),
            "demo": demo,
            "suffix": args.suffix,
            "response": response,
            "n_response_tokens": int(new_tokens.shape[0]),
            "pred_xy": pred,
            "no_pred": pred is None,
            "click_in_gt": bool(ok),
        }) + "\n")
        out_f.flush()

        if n_processed % args.batch_size_log == 0 or n_processed == len(data):
            elapsed = time.time() - batch_start_t
            cum_acc = 100.0 * n_correct / n_processed
            batch_acc = 100.0 * batch_correct / args.batch_size_log if n_processed % args.batch_size_log == 0 else 0.0
            print(f"[batch {n_processed//args.batch_size_log:3d}] "
                  f"sample {n_processed}/{len(data)}  "
                  f"batch_acc={batch_acc:.1f}% cum={cum_acc:.1f}% np={n_no_pred} "
                  f"({elapsed:.1f}s for last batch)",
                  flush=True)
            batch_start_t = time.time()
            batch_correct = 0; batch_np = 0

    out_f.close()
    acc = 100.0 * n_correct / len(data)
    print(f"\n=== Final: {n_correct}/{len(data)} = {acc:.2f}% click-accuracy ===")
    print(f"  parse failures: {n_no_pred}/{len(data)}")
    print(f"  per-sample preds: {args.out}")


if __name__ == "__main__":
    main()