#!/usr/bin/env python
"""Stage-b GRPO entry: gates -> worker -> banner -> train -> stage manifest.

    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
        configs/grpo_stageb.toml
    # tiny end-to-end smoke (2 prompts x G8 x 32 tokens x 10 steps):
    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
        configs/grpo_stageb.toml --smoke

Topology (v1, DESIGN v3 §B.1b): 3 FSDP train ranks + one dedicated vLLM
engine GPU (default 3, override NMO_ROLLOUT_GPU). Rank 0 owns all protocol
I/O; checkpoint saves ARE the weight sync (the worker reloads the engine
whenever a request's weights_dir changes). The loop is section-07's Trainer
verbatim -- GrpoBatchSource yields lazy handles and the loss-adapter wrapper
materialises each optimizer step's rollout at consumption time.

Launch gates, ordered by cost: config validation (denominator arithmetic,
cadence/sync alignment, R3 clipping, R4 canary) -> ceiling receipt
(0/100-perfect-policy trap) -> banding manifest -> topology check ->
A3 conformance + entropy-when-armed (module construction) -> horizon
consistency (scheduler vs configured steps).
"""

from __future__ import annotations

import argparse
import atexit
import json
import logging
import os
import sys
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("grpo")


#: The rollout worker runs in its OWN session (torchrun's teardown does not reach it) and used to
#: serve forever. If rank 0 is SIGKILLed it is orphaned holding ~85% of its GPU, so it now exits by
#: itself after this long with no request -- far above any step/checkpoint gap, far below "forever".
ROLLOUT_IDLE_EXIT_S = 7200.0


def _spawn_rollout_worker(spool: str, gpu: int, log_path: str):
    """Spawn the vLLM worker and tie its teardown to rank 0's exit.

    Everything between the spawn and the train try/finally (Trainer construction, the horizon
    check, canary data, resume_from_latest) can raise; atexit covers every uncaught-exception path.
    `terminate_worker` is idempotent, so the normal finally and this hook do not conflict.
    """
    from align.rollout import client as rollout_client

    proc = rollout_client.spawn_worker(spool, gpu=gpu, log_path=log_path, idle_exit_s=ROLLOUT_IDLE_EXIT_S)
    atexit.register(rollout_client.terminate_worker, proc)
    return proc


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("config")
    ap.add_argument("--smoke", action="store_true", help="tiny shapes, 10 steps")
    ap.add_argument("--band-file", default=None, help="override the configured band file")
    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
    ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
    args = ap.parse_args()

    apply_resource_guards(
        cpu_threads=8,
        heavy=args.heavy_lock,
        label="grpo-stageb",
        strict_host_check=args.strict_host_check,
    )

    from align.utils.dist import init_distributed
    from align.rl.stage_config import load_stage_config as _early_stage_config

    # Run dir from the config (parsed before the process group: a bad TOML fails
    # in milliseconds, and the NCCL flight-recorder dump lands under it).
    init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir)

    from align.train.stage_gates import assert_distinct_devices

    assert_distinct_devices()

    import torch
    import torch.distributed as dist

    from align.data.schema import build_prompt, extract_answer, register_answer_suffixes
    from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test
    from align.eval.verifier import OutcomeVerifier, VerifierConfig
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.rl.batch_source import GrpoBatchSource
    from align.rl.canary import TaskAccuracyCanary
    from align.rl.grpo_module import GRPOTrainModule
    from align.rl.rewards import CompositeReward
    from align.rl.stage_config import load_stage_config
    from align.rl.tripwire import CollapseTripwire
    from align.rollout.client import make_rollout_client
    from align.rollout.protocol import RolloutRequest
    from align.train.dpo_stage import resolve_sft_checkpoint
    from align.train.trainer import Trainer
    from align.utils.guard import assert_paths_off_root

    rank = dist.get_rank() if dist.is_initialized() else 0
    world = dist.get_world_size() if dist.is_initialized() else 1

    cfg = load_stage_config(args.config, smoke=args.smoke)
    if args.band_file:
        cfg.data.band_file = args.band_file
        cfg.validate()
    assert_paths_off_root([cfg.run_dir])
    ceiling = cfg.validate_ceiling_receipt()
    cfg.validate_world(world)
    with open(cfg.data.banding_manifest) as f:
        banding_manifest = json.load(f)

    register_answer_suffixes("</s>")  # BEFORE any extraction (golden.py's incident)

    # -- data: the band's prompt rows ---------------------------------------
    prompt_rows = []
    with open(cfg.data.band_file) as f:
        for line in f:
            if line.strip():
                prompt_rows.append(json.loads(line))
    logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows))

    # -- policy + FSDP -------------------------------------------------------
    sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None)
    logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source)
    from align.model.hf_config import read_config
    from align.tokenizer import load_tokenizer

    tokenizer = load_tokenizer(sft_checkpoint)
    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

    model_cfg = read_config(sft_checkpoint)
    model = NmoMistralForCausalLM.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)),
        attn_implementation=cfg.train.attn_implementation,
        device="cuda",
    ).to("cuda", dtype=torch.bfloat16)

    from align.train.fsdp import apply_activation_checkpointing, apply_fsdp

    dist_model = apply_fsdp(apply_activation_checkpointing(model))
    logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30)

    # -- rollout backend (rank 0 owns it) ------------------------------------
    # 'spool': external vLLM worker process on NMO_ROLLOUT_GPU (own venv).
    # 'native': in-repo torch sampler on a STANDALONE engine copy living on that
    # same dedicated GPU, in rank 0's process. It must NOT be the FSDP2-sharded
    # model: decode forwards are collectives the other ranks would never join
    # (deadlock). Either way the GPU is outside the training ranks' devices.
    rollout_gpu = int(os.environ.get("NMO_ROLLOUT_GPU", "3"))
    spool = os.path.join(cfg.run_dir, "spool")
    worker_proc = None
    client = None
    if cfg.rollout.backend == "spool":
        if rank == 0:
            worker_proc = _spawn_rollout_worker(
                spool, rollout_gpu, os.path.join(cfg.run_dir, "worker.log")
            )
            logger.info("rollout worker spawned (pid %d) on GPU %s", worker_proc.pid, rollout_gpu)
        client = make_rollout_client("spool", spool=spool, timeout_s=cfg.rollout_timeout_s)
    elif rank == 0:
        client = _build_native_client(
            model_cfg, sft_checkpoint, tokenizer, rollout_gpu, world, cfg.rollout.max_response_len
        )
        logger.info("native rollout engine on cuda:%d (standalone copy of %s)", rollout_gpu, sft_checkpoint)

    # -- module + trainer -----------------------------------------------------
    from align.train.offload_adamw import make_optimizer

    run_cfg = _run_config(cfg)
    optimizer = make_optimizer(run_cfg, [p for p in dist_model.parameters() if p.requires_grad])
    collapse = CollapseTripwire()
    module = GRPOTrainModule(
        dist_model,
        optimizer,
        grad_accum=cfg.train.grad_accum,
        grpo_cfg=cfg.grpo,
        temperature=cfg.rollout.temperature,
        collapse_tripwire=collapse,
        max_grad_norm=cfg.train.max_grad_norm,
    )

    source = GrpoBatchSource(
        prompt_rows,
        prompts_per_step=cfg.rollout.prompts_per_step,
        n_samples_per_prompt=cfg.grpo.num_samples_per_prompt_rollout,
        max_response_len=cfg.rollout.max_response_len,
        temperature=cfg.rollout.temperature,
        optimizer_steps=cfg.train.optimizer_steps,
        grad_accum=cfg.train.grad_accum,
        seed=cfg.train.seed,
        prompt_builder=build_prompt,
        dead_group_retries=cfg.dead_group_retries,
    )
    source.bind_pad(pad_id)

    trainer = Trainer(
        run_cfg,
        dist_model,
        source,
        module._grpo_loss_adapter,
        run_dir=cfg.run_dir,
        stage="grpo",
        train_module=module,
    )
    # HORIZON CONSISTENCY (the Arm A zombie-LR class): the scheduler's total
    # must be exactly the configured optimizer_steps -- this loop generates
    # its own data, so any mismatch is arithmetic, not sampling.
    if trainer.scheduler.total_steps != cfg.train.optimizer_steps:
        raise RuntimeError(
            f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs "
            f"configured optimizer_steps={cfg.train.optimizer_steps}"
        )

    # -- canary (R4, rank 0) --------------------------------------------------
    canary = None
    canary_generate = None
    if rank == 0:
        sources = DataSources.from_local_hub()
        gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems]
        canary_problems = [
            {"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows
        ]

        def canary_generate(prompts: list[str]) -> list[str]:
            req = RolloutRequest(
                request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}",
                weights_dir=_latest_weights(trainer, sft_checkpoint),
                prompts=prompts,
                n_samples_per_prompt=1,
                max_tokens=cfg.rollout.max_response_len,
                temperature=0.0,
            )
            client.submit(req)
            result = client.collect(req.request_id)
            return [texts[0] for texts in result.texts]

        canary = TaskAccuracyCanary(
            problems=canary_problems,
            generate_fn=canary_generate,
            verifier=OutcomeVerifier(VerifierConfig(workers=4)),
            extract_fn=extract_answer,
            every=cfg.canary.every,
        )

    source.bind(
        client=client,
        reward=CompositeReward(
            format_weight=cfg.reward.format_weight,
            overlong_max_len=cfg.rollout.max_response_len,
            overlong_cache_len=cfg.rollout.overlong_cache_len,
        ),
        verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None,
        old_logprob_fn=module.compute_logprobs_nograd,
        weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint),
        rank=rank,
        world_size=world,
        device="cuda",
        canary=canary,
        collapse_tripwire=collapse,
        trainer=trainer,
        canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"),
    )
    trainer.loss_adapter = source.lazy_adapter(module)

    # -- resume ---------------------------------------------------------------
    # A relaunch continues the newest checkpoint in run_dir, if any. Before
    # this was wired, a relaunch restarted from the SFT substrate. What a
    # resume restores, and where each piece lives:
    #   weights + optimizer + schedule + stream offset   Trainer (resume.pt)
    #   CollapseTripwire windows and latch                module extra_state
    #   prompt order, cursor, dead-group stats, alarms    source state_dict
    #   rollout engine weights                            _latest_weights: the
    #       Trainer registers the resumed checkpoint, so the first post-resume
    #       rollout reloads from it instead of the substrate.
    resumed_from = trainer.resume_from_latest()
    if resumed_from is not None:
        logger.info(
            "RESUMED from %s at optimizer step %d / micro step %d",
            resumed_from,
            trainer.optimizer_step,
            trainer.micro_step,
        )

    banner = {
        "stage": "grpo_stageb" + ("_smoke" if args.smoke else ""),
        "substrate": sft_checkpoint,
        "resumed_from": resumed_from,
        "band_file": cfg.data.band_file,
        "band_prompts": len(prompt_rows),
        "shape": {
            "prompts_per_step": cfg.rollout.prompts_per_step,
            "G": cfg.grpo.num_samples_per_prompt_rollout,
            "max_response_len": cfg.rollout.max_response_len,
            "loss_denominator": cfg.grpo.loss_denominator,
            "optimizer_steps": cfg.train.optimizer_steps,
        },
        "ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()},
        "banding_manifest_rows": banding_manifest.get("rows"),
        "topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")},
        "rectifications": (
            f"R3 clip={cfg.train.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, "
            "R5 pin-on-fire, reduction-time denominator assert armed"
        ),
    }
    logger.info("BANNER %s", json.dumps(banner, indent=1))

    # -- train ----------------------------------------------------------------
    try:
        summary = trainer.train()
        checkpoint = trainer.save_checkpoint()
    finally:
        if worker_proc is not None:
            from align.rollout.client import terminate_worker

            terminate_worker(worker_proc)
    if trainer.stop_exit_code:
        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
        # optimizer-step boundary; the stage is NOT finished. Exit with the
        # documented status instead of evaluating / reporting a half-trained arm.
        raise SystemExit(trainer.stop_exit_code)

    result = {
        "stage": banner["stage"],
        "summary": summary,
        "checkpoint": checkpoint,
        "banner": banner,
        "canary_history": canary.history if canary else [],
        "collapse_alarms": source.alarms,
        "collapse_tripwire": collapse.as_log_dict(),
        "dead_groups": source.filter_stats.as_log_dict(),
        "dropped_remainder_rows": source.dropped_remainder_rows,
    }
    if rank == 0:
        path = os.path.join(cfg.run_dir, "stage_result.json")
        with open(path + ".partial", "w") as f:
            json.dump(result, f, indent=1, default=str)
        os.replace(path + ".partial", path)
        logger.info("stage result: %s", path)
    return 0


def _build_native_client(
    model_cfg, sft_checkpoint: str, tokenizer, rollout_gpu: int, world: int, max_response_len: int
):
    """Rank 0's native rollout backend on a dedicated, standalone engine copy.

    The copy is built from the SFT substrate (the policy's step-0 weights) and
    refreshed from the latest saved checkpoint whenever `weights_dir` changes
    (`checkpoint_weight_loader`). Refuses a device a training rank already
    owns: training ranks are bound to local_rank 0..world-1, so the engine GPU
    index must be >= world, and visible.
    """
    import torch

    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.rollout import checkpoint_weight_loader, make_rollout_client
    from align.rollout.native import NativeRolloutConfig

    if rollout_gpu < world or rollout_gpu >= torch.cuda.device_count():
        raise RuntimeError(
            f"NMO_ROLLOUT_GPU={rollout_gpu} is not a free visible device: training ranks own "
            f"cuda:0..{world - 1} and {torch.cuda.device_count()} device(s) are visible. The "
            "engine copy must live on its own GPU (launch with all four GPUs visible)."
        )
    device = torch.device("cuda", rollout_gpu)
    engine_model = (
        NmoMistralForCausalLM.from_state_dict(
            model_cfg,
            ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint, dtype=torch.bfloat16)),
            device=str(device),
        )
        .to(device, dtype=torch.bfloat16)
        .eval()
    )
    return make_rollout_client(
        "native",
        model=engine_model,
        tokenizer=tokenizer,
        native_config=NativeRolloutConfig(
            max_seq_len_cap=2048,
            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0,
        ),
        load_weights_fn=checkpoint_weight_loader(engine_model),
        device=device,
    )


def _latest_weights(trainer, sft_checkpoint: str) -> str:
    """The engine's weights source: latest saved checkpoint, else the substrate."""
    if trainer.checkpoints_written:
        return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}")
    return sft_checkpoint


def _run_config(cfg):
    from align.config import RunConfig

    run = RunConfig()
    run.optimizer.name = cfg.train.optimizer_name
    run.optimizer.lr = cfg.train.lr
    run.optimizer.grad_accum = cfg.train.grad_accum
    run.train.epochs = 1  # the horizon is optimizer_steps; the source cycles prompts itself
    run.train.micro_batch_size = cfg.train.micro_batch_size
    run.train.save_every_micro_steps = cfg.train.save_every_micro_steps
    run.train.keep_checkpoints = cfg.train.keep_checkpoints
    run.train.per_rank_optimizer_state = cfg.train.per_rank_optimizer_state
    run.train.max_wall_seconds = cfg.train.max_wall_seconds
    run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps
    run.train.max_skipped_fraction = cfg.train.max_skipped_fraction
    run.train.warmup_frac = cfg.train.warmup_frac
    run.train.full_ft = True
    run.train.seed = cfg.train.seed
    run.tripwires.armed = tuple(cfg.armed_tripwires)
    run.run_name = "grpo_stageb"
    return run


if __name__ == "__main__":
    from align.exit_codes import run_with_exit_codes

    raise SystemExit(run_with_exit_codes(main))