"""Zero-shot localization of the runic inscription with Qwen3.5 (bounding box in 0-1000 coordinates).

Writes one JSON line per image: {"image", "raw", "box"} where box is [x1, y1, x2, y2] in pixels of
the stored image, or null when no box could be parsed. Resumable.
"""
import argparse
import json
import os
import re
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from qwen_common import clean_generation, load_model, stop_ids  # noqa: E402
from PIL import Image

PROMPT = ("Locate the runic inscription in this image: the area that contains the carved or painted runes, "
          "including all lines of runes. Output only its bounding box as JSON: "
          "{\"bbox_2d\": [x1, y1, x2, y2]} with coordinates between 0 and 1000.")


def parse_box(text, w, h):
    m = re.search(r"\[([^\[\]]*)\]", text)  # first coordinate list; the "2" in the key "bbox_2d" is not a coordinate
    nums = re.findall(r"-?\d+(?:\.\d+)?", m.group(1) if m else "")
    if len(nums) < 4:
        return None
    x1, y1, x2, y2 = [float(v) for v in nums[:4]]
    if max(x1, y1, x2, y2) <= 1.0:  # relative 0-1
        scale_x, scale_y = w, h
    else:
        scale_x, scale_y = w / 1000.0, h / 1000.0
    x1, x2 = sorted((x1 * scale_x, x2 * scale_x))
    y1, y2 = sorted((y1 * scale_y, y2 * scale_y))
    x1, y1 = max(0, x1), max(0, y1)
    x2, y2 = min(w, x2), min(h, y2)
    if x2 - x1 < 0.05 * w or y2 - y1 < 0.05 * h:
        return None
    return [round(x1), round(y1), round(x2), round(y2)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--inputs", required=True)
    ap.add_argument("--image_root", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--bs", type=int, default=16)
    args = ap.parse_args()
    images = []
    for p in args.inputs.split(","):
        for l in open(p):
            im = json.loads(l)["image"]
            if im not in images:
                images.append(im)
    done = set()
    if os.path.exists(args.out):
        done = {json.loads(l)["image"] for l in open(args.out)}
    todo = [i for i in images if i not in done]
    print("images", len(images), "todo", len(todo), flush=True)
    processor, model = load_model(args.model)
    model.to("cuda").eval()
    tok = processor.tokenizer
    tok.padding_side = "left"
    msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": PROMPT}]}]
    prompt = processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    fout = open(args.out, "a")
    for b in range(0, len(todo), args.bs):
        chunk = todo[b:b + args.bs]
        ims = []
        for name in chunk:
            im = Image.open(os.path.join(args.image_root, name)).convert("RGB")
            ims.append(im)
        enc = processor(text=[prompt] * len(ims), images=ims, return_tensors="pt", padding=True)
        enc = {k: v.to("cuda") for k, v in enc.items()}
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            out = model.generate(**enc, max_new_tokens=40, do_sample=False, eos_token_id=stop_ids(tok),
                                 pad_token_id=tok.pad_token_id)
        gen = out[:, enc["input_ids"].shape[1]:]
        for name, im, g in zip(chunk, ims, gen):
            raw = clean_generation(tok.decode(g, skip_special_tokens=False))
            fout.write(json.dumps({"image": name, "raw": raw, "box": parse_box(raw, *im.size)}) + "\n")
        fout.flush()
        print(f"{min(b + args.bs, len(todo))}/{len(todo)}", flush=True)
    print("DONE", flush=True)


if __name__ == "__main__":
    main()