"""
Generate the synthetic Arabic math-reasoning corpus with google/gemma-4-12B-it.

Resumable by construction:

  * A "task" is one generation call that asks for ITEMS_PER_TASK problems. Task N's prompt is a
    pure function of N (synth_common.build_task), so nothing about the plan is stored — restart
    the script and it redraws the identical prompts.
  * Every completed task is appended to out_synth/generations.jsonl as one line and fsynced.
    On start the file is replayed, completed task_ids are skipped, and a torn final line (a
    machine that died mid-write) is dropped. Worst case a crash costs one in-flight batch.
  * The stop condition is *accepted rows*, not tasks: each cached generation is re-validated on
    load, so a resumed run knows how many good rows it already has and issues only what's missing.

Backends:
  BACKEND=hf     transformers batched generation                       (default, works here)
  BACKEND=vllm   vLLM continuous batching                             (see the note below)

vLLM cannot serve Gemma 4 on this box: every vLLM release that knows the gemma4 architecture
(>= 0.20.2) pins torch 2.11, which is a CUDA 13 build, and this host's driver (550 / CUDA 12.4)
caps at CUDA 12.x — vllm imports die on `libcudart.so.13`. The vLLM path is kept working for a
box with driver >= 580, or for a gemma3/Qwen generator on the pinned .venv-vllm stack.

Usage:
  P=/notebooks/50M/.venv-lfm2/bin/python
  TARGET=100000 $P -u synth_generate.py
"""
import json
import os
import sys
import time
from pathlib import Path

import synth_common as sc

MODEL_DIR   = os.environ.get("GEN_MODEL", "./models/gemma-4-12B-it")
OUT_DIR     = Path(os.environ.get("OUT_DIR", "out_synth"))
CACHE       = OUT_DIR / "generations.jsonl"
TARGET      = int(os.environ.get("TARGET", 100_000))
BATCH       = int(os.environ.get("BATCH", 32))
MAX_NEW     = int(os.environ.get("MAX_NEW", 1400))
TEMPERATURE = float(os.environ.get("TEMPERATURE", 0.9))
TOP_P       = float(os.environ.get("TOP_P", 0.95))
SEED        = int(os.environ.get("SEED", 1234))
BACKEND     = os.environ.get("BACKEND", "hf")
# Recorded on every cached row so a merged multi-node corpus stays attributable.
MODEL_NAME  = os.environ.get("GEN_MODEL_NAME", os.path.basename(MODEL_DIR.rstrip("/")))
# Multimodal wrappers can default to eager attention, which is several times slower to decode.
ATTN        = os.environ.get("ATTN", "sdpa")
MAX_TASKS   = int(os.environ.get("MAX_TASKS", 400_000))
# Two machines generating at once must own disjoint task-id ranges: build_task() is a pure
# function of the id, so overlapping ranges redraw byte-identical prompts and every row the
# second machine produces dies as a duplicate template in build_synth_dataset.py. START_TASK
# offsets this node's range; MAX_TASKS is counted from there, not from zero.
START_TASK  = int(os.environ.get("START_TASK", 0))
# Which slice of the variation grid to draw from. "default" is the original 15-operation grid;
# "relational" is the comparison pool (x = 2y, x = y/2, x = y + n ...) that the default grid never
# produced. Use a disjoint START_TASK for a relational run — same rule as two nodes not colliding.
POOL        = os.environ.get("POOL", "default")
# Qwen3-style templates open a <think> block in the generation prompt unless this is passed,
# and the model then spends the whole MAX_NEW budget reasoning before it ever emits a tag.
CHAT_KWARGS = {"enable_thinking": False} if os.environ.get("NO_THINK", "0") == "1" else {}


def load_cache():
    """-> (set of finished task_ids, accepted row count). Tolerates a truncated final line."""
    done, accepted = set(), 0
    if not CACHE.exists():
        return done, accepted
    with open(CACHE, encoding="utf-8") as fh:
        for line in fh:
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                print("[!] dropping truncated final line of the cache")
                continue
            done.add(rec["task_id"])
            for item in sc.parse_items(rec["raw"]):
                ok, _ = sc.validate(item)
                accepted += ok
    return done, accepted


class HFBackend:
    """Batched transformers generation. No paged attention, but B sequences decode in parallel."""

    def __init__(self):
        import torch
        import transformers
        from transformers import AutoProcessor, AutoTokenizer

        self.torch = torch
        try:
            self.tok = AutoProcessor.from_pretrained(MODEL_DIR).tokenizer
        except Exception:
            self.tok = AutoTokenizer.from_pretrained(MODEL_DIR)
        self.tok.padding_side = "left"
        if self.tok.pad_token_id is None:
            self.tok.pad_token = self.tok.eos_token
        print(f"[*] loading {MODEL_DIR} (bf16)", flush=True)
        # gemma-4-*-it is a unified multimodal checkpoint, so the plain causal-LM auto class does
        # not always claim it. Try the multimodal auto classes first and fall back.
        self.model, last = None, None
        for name in ("AutoModelForMultimodalLM", "AutoModelForImageTextToText",
                     "AutoModelForCausalLM"):
            cls = getattr(transformers, name, None)
            if cls is None:
                continue
            try:
                self.model = cls.from_pretrained(
                    MODEL_DIR, dtype=torch.bfloat16, device_map="cuda:0",
                    attn_implementation=ATTN).eval()
                print(f"    loaded via {name}", flush=True)
                break
            except Exception as exc:                       # noqa: BLE001 - report the last failure
                last = f"{name}: {exc}"
        if self.model is None:
            raise RuntimeError(f"could not load {MODEL_DIR}; last error -> {last}")
        self.model.config.use_cache = True

    def generate(self, prompts):
        texts = [self.tok.apply_chat_template([{"role": "user", "content": p}],
                                              tokenize=False, add_generation_prompt=True,
                                              **CHAT_KWARGS)
                 for p in prompts]
        enc = self.tok(texts, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda:0")
        with self.torch.no_grad():
            out = self.model.generate(**enc, max_new_tokens=MAX_NEW, do_sample=True,
                                      temperature=TEMPERATURE, top_p=TOP_P,
                                      pad_token_id=self.tok.pad_token_id)
        width = enc["input_ids"].shape[1]
        return [self.tok.decode(seq[width:], skip_special_tokens=True) for seq in out]


class VLLMBackend:
    def __init__(self):
        from vllm import LLM, SamplingParams

        kw = {}
        if os.environ.get("QUANT"):                 # e.g. QUANT=modelopt for NVFP4 checkpoints
            kw["quantization"] = os.environ["QUANT"]
        self.llm = LLM(model=MODEL_DIR, dtype=os.environ.get("DTYPE", "bfloat16"),
                       max_num_seqs=BATCH,
                       gpu_memory_utilization=float(os.environ.get("GPU_UTIL", 0.90)),
                       max_model_len=int(os.environ.get("MAX_LEN", 4096)), **kw)
        self.params = SamplingParams(temperature=TEMPERATURE, top_p=TOP_P, max_tokens=MAX_NEW,
                                     seed=None)
        self.tok = self.llm.get_tokenizer()

    def generate(self, prompts):
        texts = [self.tok.apply_chat_template([{"role": "user", "content": p}],
                                              tokenize=False, add_generation_prompt=True,
                                              **CHAT_KWARGS)
                 for p in prompts]
        outs = self.llm.generate(texts, self.params)
        return [o.outputs[0].text for o in outs]


def main():
    OUT_DIR.mkdir(exist_ok=True)
    done, accepted = load_cache()
    print(f"[*] cache: {len(done):,} tasks done, {accepted:,} rows accepted "
          f"(target {TARGET:,}) | model {MODEL_NAME} | tasks {START_TASK:,}.."
          f"{START_TASK + MAX_TASKS:,}", flush=True)
    if accepted >= TARGET:
        print("[+] target already met — nothing to do")
        return

    backend = (VLLMBackend if BACKEND == "vllm" else HFBackend)()

    next_id = START_TASK
    last_id = START_TASK + MAX_TASKS
    started, gen_rows, gen_tasks = time.time(), 0, 0
    fh = open(CACHE, "a", encoding="utf-8")
    while accepted < TARGET and next_id < last_id:
        batch = []
        while len(batch) < BATCH and next_id < last_id:
            if next_id not in done:
                batch.append((next_id, *sc.build_task(next_id, SEED, POOL)))
            next_id += 1
        if not batch:
            break

        t0 = time.time()
        raws = backend.generate([p for _, _, p in batch])
        batch_accept = 0
        for (task_id, axes, _), raw in zip(batch, raws):
            fh.write(json.dumps({"task_id": task_id, "axes": axes, "raw": raw,
                                 "model": MODEL_NAME}, ensure_ascii=False) + "\n")
            for item in sc.parse_items(raw):
                ok, _ = sc.validate(item)
                batch_accept += ok
        fh.flush()
        os.fsync(fh.fileno())          # a crash costs the in-flight batch, never the cache

        accepted += batch_accept
        gen_rows += batch_accept
        gen_tasks += len(batch)
        dt = time.time() - t0
        rate = gen_rows / max(time.time() - started, 1e-9)
        eta = (TARGET - accepted) / rate / 3600 if rate > 0 else float("inf")
        print(f"[{accepted:>7,}/{TARGET:,}] +{batch_accept:>3} rows  "
              f"batch {len(batch)} in {dt:5.1f}s  "
              f"accept {batch_accept / (len(batch) * sc.ITEMS_PER_TASK):5.1%}  "
              f"{rate * 3600:,.0f} rows/h  ETA {eta:4.1f}h", flush=True)

    fh.close()
    print(f"[+] {accepted:,} accepted rows in cache after {gen_tasks:,} new tasks")


if __name__ == "__main__":
    sys.exit(main())