"""Greedy generation with Qwen3.5 (+ optional LoRA adapter) for recognition or translation.

Input jsonl rows:
  ocr: {"id", "image", ...}              -> writes {"id", "pred", ...}
  mt : {"id", <src_field>, "lang", ...}  -> writes {"id", "pred", ...}
Resumable: ids already present in --out are skipped. Rows are written after every batch.
"""
import argparse
import json
import os
import sys

import torch

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", choices=["mt", "ocr"], required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--adapter", default="")
    ap.add_argument("--input", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--src_field", default="src")
    ap.add_argument("--image_root", default="")
    ap.add_argument("--bs", type=int, default=16)
    ap.add_argument("--max_new_tokens", type=int, default=160)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--tags", action="store_true", help="mt: prompt with <runic-language> <2target> tags (field rl)")
    ap.add_argument("--force_rl", default="", help="with --tags: use this runic-language tag for every row")
    args = ap.parse_args()
    print("ARGS", vars(args), flush=True)

    rows = [json.loads(l) for l in open(args.input) if l.strip()]
    if args.limit:
        rows = rows[: args.limit]
    done = set()
    if os.path.exists(args.out):
        done = {json.loads(l)["id"] for l in open(args.out) if l.strip()}
    todo = [r for r in rows if r["id"] not in done]
    print(f"rows {len(rows)} done {len(done)} todo {len(todo)}", flush=True)
    if not todo:
        return

    processor, model = load_model(args.model, adapter=args.adapter or None)
    model.to("cuda").eval()
    tok = processor.tokenizer
    tok.padding_side = "left"
    eos = stop_ids(tok)
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    fout = open(args.out, "a")
    for b in range(0, len(todo), args.bs):
        chunk = todo[b : b + args.bs]
        if args.task == "ocr":
            texts = [prompt_text(processor, "ocr") for _ in chunk]
            images = [load_image(os.path.join(args.image_root, r["image"])) for r in chunk]
            enc = processor(text=texts, images=images, return_tensors="pt", padding=True)
        else:
            texts = [prompt_text(processor, "mt", r[args.src_field] or "", r["lang"],
                                 (args.force_rl or r["rl"]) if args.tags else None) for r in chunk]
            enc = processor(text=texts, return_tensors="pt", padding=True, **encode_kwargs(processor))
        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=args.max_new_tokens, do_sample=False,
                                 eos_token_id=eos, pad_token_id=tok.pad_token_id, repetition_penalty=1.0)
        gen = out[:, enc["input_ids"].shape[1]:]
        for r, g in zip(chunk, gen):
            pred = clean_generation(tok.decode(g, skip_special_tokens=False))
            rec = dict(r)
            rec["pred"] = pred
            fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
        fout.flush()
        print(f"{min(b + args.bs, len(todo))}/{len(todo)}", flush=True)
    print("DONE", flush=True)


if __name__ == "__main__":
    main()