#!/usr/bin/env python
"""Generate completions with the align engine (or the bucketed fallback).

    uv run python scripts/generate.py \
        --checkpoint /apps/development/rlhf/nano-mistral-o1/checkpoints/sft \
        --prompts data/eval/math500.json --out .tmp/gen/math500.jsonl

Sharding across the 4 A10Gs is done by the CALLER, one process per device --
the engine is single-replica and knows nothing about other GPUs (no
collectives, no NCCL). The idiom, which respects the one-heavy-job rule:

    for i in 0 1 2 3; do
      RANK=$i CUDA_VISIBLE_DEVICES=$i uv run python scripts/generate.py \
        --shard-index $i --num-shards 4 --out .tmp/gen/shard$i.jsonl &
    done; wait

`acquire_heavy_job_lock` is rank-aware, so only RANK=0 takes the exclusive
flock; shards 1.. do NOT take it, they VERIFY it (`verify_shard_lock`): a lock
held by anything other than another generate.py (their own shard 0) refuses
the launch, so a live training stage cannot be overlapped by three of the four
shards. Every shard also runs the host preflight. Never launch this against a
live training stage: that overlap is what killed the predecessor's first SFT
run. Completions carry rid "<shard>:<local i>"; the global prompt index of
local i is shard_index + i * num_shards (stride sharding).

Every output path is asserted onto /apps (`/` has ~13GB and /tmp is a noexec
tmpfs). The launch banner prints the checkpoint fingerprint, the attention
backend that ACTUALLY ran, the engine budget, the sampling params and the
per-gate status; a missing mandatory key raises rather than printing a partial
banner, and gate status is "PENDING" until recorded -- never assumed passed.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.guard import apply_resource_guards  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("generate")

PROJECT_ROOT = "/apps/development/rlhf/nano-mistral-o1"


def validate_shards(shard_index: int, num_shards: int) -> None:
    """An out-of-range shard silently yields an EMPTY prompt list (and a "successful" empty run)."""
    if num_shards < 1 or not (0 <= shard_index < num_shards):
        raise ValueError(f"bad shard spec: --shard-index {shard_index} must be in [0, --num-shards {num_shards})")


def _proc_cmdline(pid: int) -> str:
    try:
        with open(f"/proc/{pid}/cmdline", "rb") as f:
            return f.read().replace(b"\0", b" ").decode(errors="replace")
    except OSError:
        return ""


def verify_shard_lock(rank: int, lock_inode: int | None = None, locks_text: str | None = None, cmdline_of=_proc_cmdline) -> None:
    """Shards 1.. never take the heavy-job flock (`acquire_heavy_job_lock` returns True for RANK != 0),
    so without this a generation launched while a trainer holds it would be refused on shard 0 only,
    while shards 1-3 each loaded a 7B onto GPUs 1-3. Read-only: scans /proc/locks for the holder of the
    lock file (taking the flock here would race shard 0's own non-blocking acquire). The only holder
    tolerated is another generate.py -- this launch's own shard 0."""
    if rank == 0:
        return
    from align.exit_codes import PreflightRefused
    from align.utils.guard import LOCK_PATH

    if lock_inode is None:
        try:
            lock_inode = os.stat(LOCK_PATH).st_ino
        except FileNotFoundError:
            return  # no lock file yet, so no holder
    if locks_text is None:
        with open("/proc/locks") as f:
            locks_text = f.read()
    for line in locks_text.splitlines():
        parts = line.split()
        if len(parts) < 6 or parts[1] != "FLOCK" or parts[5].split(":")[-1] != str(lock_inode):
            continue
        holder = cmdline_of(int(parts[4]))
        if "generate.py" not in holder:
            raise PreflightRefused(
                f"another heavy align job holds the lock (pid {parts[4]}: {holder.strip()[:120]!r}); this shard "
                f"(RANK={rank}) must not overlap it. Shard 0 refuses for the same reason."
            )


def read_prompts(path: str) -> list[str]:
    """Accept .json (list of str, or list of dicts with a prompt/problem/text
    key), .jsonl (same per line), or .txt (one prompt per line). Refuses an
    empty file loudly: an eval that silently generates nothing looks like a
    fast eval."""

    def _field(obj):
        if isinstance(obj, str):
            return obj
        for k in ("prompt", "problem", "question", "text"):
            if k in obj:
                return obj[k]
        raise ValueError(
            f"prompt record has none of prompt/problem/question/text: {sorted(obj)[:8]}"
        )

    if path.endswith(".jsonl"):
        with open(path) as f:
            rows = [_field(json.loads(line)) for line in f if line.strip()]
    elif path.endswith(".json"):
        with open(path) as f:
            rows = [_field(o) for o in json.load(f)]
    else:
        with open(path) as f:
            rows = [line.rstrip("\n") for line in f if line.strip()]
    if not rows:
        raise ValueError(f"{path} contains no prompts")
    return rows


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--checkpoint", required=True, help="HF-layout checkpoint directory (safetensors)"
    )
    ap.add_argument("--prompts", required=True)
    ap.add_argument("--out", required=True, help="completions JSONL; under /apps")
    ap.add_argument("--tokenizer", default=None, help="defaults to --checkpoint")
    ap.add_argument("--n-slots", type=int, default=32)
    ap.add_argument("--max-seq-len", type=int, default=1024)
    ap.add_argument("--max-new-tokens", type=int, default=300)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--top-p", type=float, default=1.0)
    ap.add_argument(
        "--stop", action="append", default=[], help="repeatable; matched on DECODED TEXT"
    )
    ap.add_argument("--stop-policy", default="trim", choices=("trim", "keep"))
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument(
        "--attn", default=None, help="force an attention backend; unavailable => hard error"
    )
    ap.add_argument(
        "--bucketed",
        action="store_true",
        help="use the static-batching fallback instead of the "
        "engine (correctness reference; slower)",
    )
    ap.add_argument(
        "--benchmark",
        action="store_true",
        help="run BOTH and write the engine-vs-static-batching speedup receipt beside --out",
    )
    ap.add_argument(
        "--hash-weights",
        action="store_true",
        help="full content hash of the checkpoint (reads 14.5GB)",
    )
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--shard-index", type=int, default=0)
    ap.add_argument("--num-shards", type=int, default=1)
    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
    ap.add_argument(
        "--heavy-lock",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="take the exclusive one-heavy-job flock. Keep it ON "
        "for any real checkpoint. --no-heavy-lock exists for "
        "the tiny-model CLI smoke test, which loads ~200KB "
        "and is not a heavy job; using it to squeeze a 7B "
        "run past a live training stage is the exact overlap "
        "that killed the predecessor's first SFT run.",
    )
    args = ap.parse_args()
    try:
        validate_shards(args.shard_index, args.num_shards)
    except ValueError as e:
        ap.error(str(e))

    apply_resource_guards(
        cpu_threads=8,
        heavy=args.heavy_lock,
        label="generate",
        strict_host_check=args.strict_host_check,
    )
    if args.heavy_lock:
        verify_shard_lock(int(os.environ.get("RANK", "0")))

    import torch

    from align.inference.bucketed import bucketed_generate
    from align.inference.engine import (
        Engine,
        GenerationRequest,
        InferenceConfig,
        benchmark_engine_vs_bucketed,
    )
    from align.inference.receipts import (
        checkpoint_fingerprint,
        engine_banner,
        load_gate_status,
        sha256_file,
        write_completions_jsonl,
    )
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.registry import render_banner
    from align.utils.guard import assert_paths_off_root

    assert_paths_off_root([os.path.dirname(os.path.abspath(args.out))])

    prompts = read_prompts(args.prompts)
    if args.limit:
        prompts = prompts[: args.limit]
    # Shard by stride, not by contiguous block: prompt files are often sorted
    # by length or difficulty, and contiguous blocks would give one replica
    # all the long ones and skew both the wall-clock and the benchmark.
    if args.num_shards > 1:
        prompts = prompts[args.shard_index :: args.num_shards]
        logger.info("shard %d/%d: %d prompts", args.shard_index, args.num_shards, len(prompts))

    from align.eval.golden import resolve_model_config
    from align.tokenizer import load_tokenizer

    tok = load_tokenizer(args.tokenizer or args.checkpoint)
    state = ckpt.load_hf_safetensors(args.checkpoint)
    # `read_config` (align.model.hf_config) states every architecture field
    # explicitly and REFUSES a config.json that omits one: transformers'
    # MistralConfig used to resolve the missing ones to its own defaults
    # (sliding_window=4096, rms_norm_eps=1e-6 -- neither of them v0.3), which
    # silently accepted a checkpoint that was a different model.
    model_cfg = resolve_model_config(args.checkpoint)
    model = NmoMistralForCausalLM.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(state),
        attn_implementation=args.attn,
        device="cuda" if torch.cuda.is_available() else "cpu",
    ).eval()
    if torch.cuda.is_available():
        model = model.to("cuda", dtype=torch.bfloat16)

    cfg = InferenceConfig(
        n_slots=args.n_slots,
        max_seq_len=args.max_seq_len,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        stop=tuple(args.stop),
        stop_string_policy=args.stop_policy,
        seed=args.seed,
        eos_token_id=tok.eos_token_id,
        pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id,
    )
    engine = Engine(model, cfg, tokenizer=tok)

    requests = [
        GenerationRequest(
            rid=f"{args.shard_index}:{i}", prompt_ids=tok(p, add_special_tokens=True)["input_ids"]
        )
        for i, p in enumerate(prompts)
    ]

    gate_status = load_gate_status(args.checkpoint)
    banner = engine_banner(
        engine=engine,
        checkpoint=checkpoint_fingerprint(args.checkpoint, args.hash_weights),
        gate_status=gate_status,
        prompts_sha256=sha256_file(args.prompts),
        extra={
            "n_prompts": len(requests),
            "path_bucketed": args.bucketed,
            "shard": f"{args.shard_index}/{args.num_shards}",
            "global_prompt_index": f"{args.shard_index} + i * {args.num_shards} (i = local index in rid)",
        },
    )
    print(render_banner(banner, title="align generate"))
    if not banner["engine_generated_numbers_usable"] and not args.bucketed:
        logger.warning(
            "the 3-part engine gate has NOT fully passed for this checkpoint "
            "(%s). These completions are fine for inspection; section-10 must "
            "use --bucketed until gate (c) is recorded.",
            gate_status,
        )

    t0 = time.perf_counter()
    if args.bucketed:
        results = bucketed_generate(model, requests, cfg, tokenizer=tok)
    else:
        results = engine.generate(requests)
    wall = time.perf_counter() - t0
    n_tokens = sum(r.n_generated_tokens for r in results)
    receipt = dict(banner)
    receipt["throughput"] = {
        "wall_s": round(wall, 3),
        "generated_tokens": n_tokens,
        "tokens_per_s": round(n_tokens / wall, 2) if wall > 0 else float("nan"),
    }
    write_completions_jsonl(args.out, results, receipt, tokenizer=tok)
    logger.info(
        "%d completions, %d tokens, %.1fs (%.1f tok/s)",
        len(results),
        n_tokens,
        wall,
        n_tokens / wall if wall else float("nan"),
    )

    if args.benchmark:
        bench_path = os.path.splitext(args.out)[0] + ".benchmark.json"
        rec = benchmark_engine_vs_bucketed(model, requests, cfg, tokenizer=tok, out_path=bench_path)
        logger.info(
            "engine vs static batching: measured speedup %.2fx "
            "(engine %.1f tok/s, bucketed %.1f tok/s)",
            rec["speedup"],
            rec["engine_tokens_per_s"],
            rec["bucketed_tokens_per_s"],
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())