"""End-to-end benchmark: reconstruct with One-D-Piece and score against a live vLLM server.

No staging directory -- images are reconstructed and posted in the same pass. That is what makes
full-size benchmarks practical: RSVQA-HR alone is 222,684 questions, so pre-reconstructing every
(question, path) pair would mean writing 445k JPEGs to lustre for a single run.

Two processes, because the two environments cannot share an interpreter:
  * server: .conda_vllm  (vllm 0.11.x -- 0.10.1.1 breaks GLM-4.1V, see serve_vllm.sh)
  * client: main env     (One-D-Piece + its timm)

Reconstruction is cached PER IMAGE, not per question. RSVQA-HR asks ~100 questions about each of
its 2,226 images; reconstructing per question would do 100x redundant tokenizer work and make the
tokenizer, not the LLM, the bottleneck.

Results stream to <out>.jsonl as they land, and a rerun skips keys already present, so a 20-hour
RSVQA-HR pass survives an interruption.

  # 1. start servers (in .conda_vllm)
  bash src/mllm/serve_vllm.sh $T/model/RSThinker 8000 0,1,2,3
  bash src/mllm/serve_vllm.sh $T/model/RSThinker 8001 4,5,6 0.85
  # 2. run the client (main env)
  python -m mllm.eval_e2e --bench rsvqa_lr --servers http://127.0.0.1:8000,http://127.0.0.1:8001 \
      --recon-gpu 7 --tag base
"""
import os, io, sys, json, glob, time, base64, random, argparse, collections, threading
from concurrent.futures import ThreadPoolExecutor

import os, sys
_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
for _p in (_ROOT, os.path.join(_ROOT, "src")):
    if _p not in sys.path:
        sys.path.insert(0, _p)
os.environ.setdefault("TOKLIP_VQ", "onedpiece")

import requests
from PIL import Image

ROOT = "${TOKLIP_ROOT}"
OUTDIR = f"{ROOT}/logs_geocot_pipeline/bench_e2e"

# Prompts follow GeoCoT380k's own phrasing, or a checkpoint trained on it is measured
# off-distribution: grounding is a BARE referring expression ("where is a windmill at the bottom")
# answered as <answer>[566, 793, 661, 871]</answer> on a 0-1000 grid, and VQA is the plain question
# with no "answer in a single word" suffix. LHRS-Bench has no training analogue (letter choice), so
# it keeps an explicit instruction.
MC_SUFFIX = "\nAnswer with the option's letter from the given choices directly."


# --------------------------------------------------------------------------- datasets
# Each loader yields dicts WITHOUT decoded pixels: full benchmarks are far too large to hold
# 222k PIL images, and one image serves ~100 questions anyway.

def iter_lhrs(n):
    import pyarrow.parquet as pq
    from satlink.eval.bench_score import pick
    f = sorted(glob.glob(f"{ROOT}/data/LHRS_Bench/**/*.parquet", recursive=True))[0]
    for i, r in enumerate(pq.read_table(f).to_pylist()):
        if n and i >= n:
            break
        yield {"key": f"lhrs:{i}", "img": ("bytes", r["image"]["bytes"]),
               "q": f"{r['question']}\n{r['choices']}{MC_SUFFIX}",
               "gt": pick(r["answer"], r["choices"]), "type": "mc", "choices": r["choices"]}


def _vrsbench(n, fname, typ):
    VB = f"{ROOT}/data/VRSBench"
    d = json.load(open(f"{VB}/{fname}"))
    if typ == "vqa":
        by = collections.defaultdict(list)
        for x in d:
            by[x["type"]].append(x)
        rng = random.Random(0)
        sel = []
        if n:
            per = max(1, n // len(by))
            for t, v in by.items():
                rng.shuffle(v)
                sel += v[:per]
        else:
            sel = d
    else:
        sel = d[:n] if n else d
    for i, x in enumerate(sel):
        p = f"{VB}/Images_val/{x['image_id']}"
        if not os.path.isfile(p):
            continue
        if typ == "vqa":
            yield {"key": f"vrsbench:{x.get('question_id', i)}", "img": ("path", p),
                   "q": x["question"], "gt": x["ground_truth"].strip().lower(),
                   "type": x["type"], "choices": None}
        else:
            # GeoCoT380k phrases every referring item as "where is <expr>" with a lower-cased
            # first letter ("where is the ship positioned towards the lower right within the
            # image."), while this eval file stores a bare capitalised noun phrase. Handing the
            # model the bare phrase makes it answer "Yes" instead of a box -- it does not read as
            # a localisation request at all.
            e = x["question"].strip()
            yield {"key": f"vrsbench_ref:{x.get('question_id', i)}", "img": ("path", p),
                   "q": "where is " + (e[0].lower() + e[1:] if e else e),
                   "gt": x["ground_truth"], "type": "ref", "choices": None}


def iter_rsvqa(n, which):
    if which == "lr":
        base, pfx, imgdir, ext = f"{ROOT}/data/RSVQA-LR", "LR", "Images_LR", "tif"
    else:
        base, pfx, imgdir, ext = f"{ROOT}/data/RSVQA-HR", "USGS", "Data", "tif"
    qs = [x for x in json.load(open(f"{base}/{pfx}_split_test_questions.json"))["questions"]
          if x.get("active")]
    ans = {x["id"]: x["answer"]
           for x in json.load(open(f"{base}/{pfx}_split_test_answers.json"))["answers"]
           if x.get("active")}
    if n:
        by = collections.defaultdict(list)
        for x in qs:
            by[x["type"]].append(x)
        rng = random.Random(0)
        per = max(1, n // len(by))
        qs = [x for t, v in by.items() for x in (rng.sample(v, min(per, len(v))))]
    # sort by image so the per-image reconstruction cache stays hot (HR: ~100 questions per image)
    qs.sort(key=lambda x: x["img_id"])
    for x in qs:
        a = ans.get((x.get("answers_ids") or [None])[0])
        if a is None:
            continue
        p = f"{base}/{imgdir}/{x['img_id']}.{ext}"
        if not os.path.isfile(p):
            continue
        yield {"key": f"rsvqa_{which}:{x['id']}", "img": ("path", p), "q": x["question"],
               "gt": str(a).strip().lower(), "type": x["type"], "choices": None}


LOADERS = {
    "lhrs":         lambda n: iter_lhrs(n),
    "vrsbench":     lambda n: _vrsbench(n, "VRSBench_EVAL_vqa.json", "vqa"),
    "vrsbench_ref": lambda n: _vrsbench(n, "VRSBench_EVAL_referring.json", "ref"),
    "rsvqa_lr":     lambda n: iter_rsvqa(n, "lr"),
    "rsvqa_hr":     lambda n: iter_rsvqa(n, "hr"),
}


# --------------------------------------------------------------------------- reconstruction
class Recon:
    """One-D-Piece reconstruction with a per-image LRU cache of the encoded JPEG bytes."""

    def __init__(self, gpu, max_tiles, nlatent, cache_size):
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
        from mllm.recon import load_onedpiece, reconstruct_batch
        self._rb = reconstruct_batch
        self.opd = load_onedpiece("cuda")
        self.max_tiles, self.nlatent = max_tiles, nlatent
        self.cache, self.order, self.cap = {}, collections.deque(), cache_size
        self.lock = threading.Lock()
        self.gpu_lock = threading.Lock()

    def _open(self, src):
        kind, val = src
        pil = Image.open(io.BytesIO(val)) if kind == "bytes" else Image.open(val)
        return pil.convert("RGB")

    def get(self, key, src):
        """-> (recon_jpeg_bytes, orig_jpeg_bytes). Both are resized to the SAME geometry so the
        only difference between the two paths is the 48-token bottleneck, not resolution."""
        with self.lock:
            hit = self.cache.get(key)
        if hit is not None:
            return hit
        pil = self._open(src)
        # Serialise the tokenizer forward pass. It shares GPU 0 with a vLLM instance that already
        # holds ~72%, so N worker threads each allocating crop activations at once blew past the
        # remaining headroom (54/690 items died on OOM even after dropping mem_util to 0.72).
        # One reconstruction at a time still does ~5 img/s, well above the ~3.5 img/s the LLM side
        # can consume, so this costs no throughput.
        with self.gpu_lock:
            rec = self._rb(self.opd, [pil], max_tiles=self.max_tiles,
                           nlatent=self.nlatent, device="cuda")[0]
        bufs = []
        for im in (rec, pil.resize(rec.size)):
            b = io.BytesIO()
            im.convert("RGB").save(b, format="JPEG", quality=95)
            bufs.append(b.getvalue())
        out = (bufs[0], bufs[1])
        with self.lock:
            self.cache[key] = out
            self.order.append(key)
            while len(self.order) > self.cap:
                self.cache.pop(self.order.popleft(), None)
        return out


# --------------------------------------------------------------------------- serving
def ask(server, jpg, question, max_new, timeout=600):
    b64 = base64.b64encode(jpg).decode()
    body = {"model": "rsthinker", "temperature": 0.0, "max_tokens": max_new,
            "messages": [{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
                {"type": "text", "text": question}]}]}
    r = requests.post(f"{server}/v1/chat/completions", json=body, timeout=timeout)
    r.raise_for_status()
    c = r.json()["choices"][0]
    return c["message"]["content"], c.get("finish_reason")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench", choices=sorted(LOADERS), required=True)
    ap.add_argument("--servers", required=True, help="comma-separated vLLM base URLs")
    ap.add_argument("--tag", default="base")
    ap.add_argument("--n", type=int, default=0, help="0 = the FULL test set")
    ap.add_argument("--recon-gpu", type=int, default=7)
    ap.add_argument("--max-tiles", type=int, default=12)
    ap.add_argument("--nlatent", type=int, default=48, help="tokens PER 256x256 crop")
    ap.add_argument("--max-new", type=int, default=1024)
    ap.add_argument("--concurrency", type=int, default=64)
    ap.add_argument("--cache-images", type=int, default=256)
    a = ap.parse_args()

    from satlink.eval.bench_score import final_answer, pick, iou_ok
    servers = a.servers.split(",")
    for s in servers:
        requests.get(f"{s}/v1/models", timeout=60).raise_for_status()
    print(f"servers ok: {servers}", flush=True)

    os.makedirs(OUTDIR, exist_ok=True)
    out_path = f"{OUTDIR}/{a.bench}.{a.tag}.jsonl"
    done = set()
    if os.path.isfile(out_path):
        for line in open(out_path):
            try:
                r = json.loads(line)
            except Exception:
                continue
            # only SUCCESSFUL rows count as done -- treating error rows as done would silently
            # bake a transient failure (an OOM burst, a server restart) into the final score
            if "error" not in r:
                done.add(r["k"])
        print(f"resuming: {len(done)} results already on disk", flush=True)

    recon = Recon(a.recon_gpu, a.max_tiles, a.nlatent, a.cache_images)
    fh = open(out_path, "a")
    wlock, n_done, t0 = threading.Lock(), [0], time.time()

    def score(txt, it):
        if it["type"] == "mc":
            return pick(txt, it["choices"]) == it["gt"]
        if it["type"] == "ref":
            return iou_ok(final_answer(txt), it["gt"])
        return final_answer(txt).strip(".").lower() == it["gt"]

    def work(idx_it):
        idx, it = idx_it
        srv = servers[idx % len(servers)]
        try:
            rec_b, orig_b = recon.get(it["img"][1] if it["img"][0] == "path" else it["key"], it["img"])
            row = {"k": it["key"], "type": it["type"], "gt": it["gt"]}
            for tag, jpg in (("recon", rec_b), ("orig", orig_b)):
                txt, fin = ask(srv, jpg, it["q"], a.max_new)
                row[tag] = {"ok": bool(score(txt, it)), "fin": fin,
                            "pred": final_answer(txt)[:120]}
            with wlock:
                fh.write(json.dumps(row) + "\n")
                n_done[0] += 1
                if n_done[0] % 200 == 0:
                    fh.flush()
                    r = n_done[0] / (time.time() - t0)
                    print(f"  {n_done[0]} done | {r:.1f} q/s", flush=True)
        except Exception as e:
            with wlock:
                fh.write(json.dumps({"k": it["key"], "error": repr(e)[:200]}) + "\n")

    items = ((i, it) for i, it in enumerate(LOADERS[a.bench](a.n)) if it["key"] not in done)
    with ThreadPoolExecutor(a.concurrency) as ex:
        list(ex.map(work, items))
    fh.close()
    print(f"wrote {out_path}", flush=True)
    report(out_path, a.bench, a.tag)


def report(path, bench, tag):
    res = {"recon": collections.defaultdict(lambda: [0, 0]), "orig": collections.defaultdict(lambda: [0, 0])}
    errs = fins = 0
    for line in open(path):
        r = json.loads(line)
        if "error" in r:
            errs += 1
            continue
        for t in ("recon", "orig"):
            res[t][r["type"]][0] += r[t]["ok"]
            res[t][r["type"]][1] += 1
            fins += r[t]["fin"] == "length"
    print(f"\n=== {bench} | {tag} | end-to-end ===")
    line = {}
    for t in ("recon", "orig"):
        h = sum(v[0] for v in res[t].values()); n = sum(v[1] for v in res[t].values())
        line[t] = 100 * h / max(1, n)
        print(f"  {t:6} {line[t]:5.1f}%  (n={n})")
    print(f"  recon cost: {line['recon'] - line['orig']:+.1f}   [errors {errs}, truncated {fins}]")
    if len(res["recon"]) > 1:
        print(f"\n  {'type':22}{'recon':>9}{'orig':>9}{'delta':>9}")
        for t in sorted(res["recon"], key=lambda t: -res["recon"][t][1]):
            x = 100 * res["recon"][t][0] / max(1, res["recon"][t][1])
            y = 100 * res["orig"][t][0] / max(1, res["orig"][t][1])
            print(f"  {t:22}{x:8.1f}%{y:8.1f}%{x - y:+8.1f}  (n={res['recon'][t][1]})")


if __name__ == "__main__":
    main()