#!/usr/bin/env python
"""SFT stage entry (section-11): gates -> wrap -> train -> eval -> quality gate.

    torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml --optimizer.lr=1e-5

Thin by design. Everything this drives is already built and tested elsewhere;
what lives here is ORDER, and the order is the product:

    golden_regression -> smoke_7b -> probe_20step -> rss_preflight
                      -> train -> eval -> quality_gate

IMPORT ORDER IS LOAD-BEARING. `apply_resource_guards()` sets
PYTORCH_CUDA_ALLOC_CONF, the /apps cache paths and the thread caps, and torch
reads all of them AT IMPORT. So this module imports only stdlib, the guard,
the config, and `stage_gates` (which is asserted torch-free); torch and every
align module that pulls it in are imported INSIDE main(), after the guards run.
A test enforces this by AST rather than by grep -- a literal search passes on
a file that mentions the guard in a comment and imports torch on the next line.

Resume semantics: rerunning after a kill picks up the gate ledger, skips gates
that already passed, and resumes the Trainer from the latest micro-step
checkpoint. Partial checkpoints are deleted at startup.
"""

from __future__ import annotations

import logging
import os
import sys
from pathlib import Path

# The package is not installed in the venv. Without this a bare `torchrun scripts/sft.py ...` (the launch
# line in docs/P2.md section 11 and in the config header) died with ModuleNotFoundError on all four ranks
# unless the operator remembered PYTHONPATH -- found by the P2 integration smoke loop.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.config import load_config
from align.train.stage_gates import (
    GateLedger,
    assemble_sft_banner,
    assert_build_is_current,
    assert_probe_resolved,
    eval_source_hash,
    evaluate_and_gate,
    parse_args,
    record_golden_from_report,
    require_golden_receipt,
    run_entry_gates,
)
from align.utils.guard import apply_resource_guards

logger = logging.getLogger(__name__)


def build_model(cfg):
    """Imported here, not at module scope: this pulls torch in.

    ARCH DISPATCH. `model.arch` selects the family; an unrecognised value RAISES (it must not
    fall through to the Mistral converters, which would either fail late on a Qwen3 snapshot's
    keys or load a different family into the wrong module).
    """
    if cfg.model.arch == "mistral":
        if cfg.lora.enabled:
            raise ValueError(
                "lora.enabled with model.arch='mistral': the dense-Mistral SFT path is full-FT only "
                "(LoRA Mistral runs live in scripts/ppo.py). Refusing rather than training full-FT "
                "while the config says LoRA."
            )
        from align.train.stage_smoke import build_stage_model

        return build_stage_model(cfg)
    if cfg.model.arch == "qwen3":
        from align.train.qwen3_stage import build_qwen3_stage_model

        return build_qwen3_stage_model(cfg)
    raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')")


def run_gates(cfg, ledger, banner_only: bool = False):
    """The M5 entry sequence. Refuses to proceed on a stale data build or an
    unvalidated eval stack before it touches a GPU.

    Non-Mistral arches have their own entry (see `qwen3_entry_banner`: the golden/MATH/GSM8K
    machinery certifies a Mistral checkpoint and does not apply)."""
    if cfg.model.arch == "qwen3":
        from align.train.qwen3_stage import qwen3_entry_banner

        return qwen3_entry_banner(cfg, ledger)
    if cfg.model.arch != "mistral":
        raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')")
    assert_build_is_current(cfg.data.manifest_path, cfg)
    # Bridge the golden RUN into the ledger by re-deciding from its persisted
    # report, immediately before the gate that reads the ledger. Without this
    # a PASSING golden run still leaves golden_regression PENDING.
    record_golden_from_report(ledger, cfg=cfg)
    require_golden_receipt(ledger, cfg)
    banner = run_entry_gates(cfg, ledger)
    assert_probe_resolved(banner)
    return banner


def train(cfg, ledger, model, banner):
    """One SFT epoch through the section-07 Trainer."""
    from align.data.loader import collate_sft
    from align.train.sft_loss import make_sft_loss_adapter
    from align.train.stage_smoke import build_training_sampler
    from align.train.trainer import Trainer

    # The SAME sampler the entry gates measured the horizon from. Building a
    # second one here is how a stage ends up timing, projecting and training
    # on three different batch streams.
    sampler, dataset = build_training_sampler(cfg)

    def adapter(m, batch):
        rows = [dataset[i] for i in batch]
        collated = collate_sft(rows, pad_token_id=cfg.data.pad_token_id).to("cuda")
        return make_sft_loss_adapter(require_fused=True)(m, (collated.input_ids, collated.labels))

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    extras = {}
    if cfg.model.arch == "qwen3":
        from align.train.qwen3_stage import qwen3_trainer_extras

        extras = qwen3_trainer_extras(cfg, model, adapter, run_dir)
    trainer = Trainer(
        cfg=cfg,
        # LoRA: the trunk-movement probe must watch the adapter's B matrices (the frozen base can
        # never move and A provably does not on step one), so it gets a view, not the raw model.
        model=extras.pop("probe_model", model),
        sampler=sampler,
        loss_adapter=adapter,
        run_dir=run_dir,
        stage="sft",
        **extras,
    )
    trainer.resume_from_latest()
    summary = trainer.train()
    if cfg.model.arch == "mistral":
        record_train_gate(ledger, trainer, summary)
    else:
        # The ledger's fixed sequence starts at the Mistral golden regression; a non-Mistral
        # stage has not passed those gates, so recording "train" would raise (or be a lie).
        logger.info("train finished (arch=%s): optimizer_step=%s", cfg.model.arch, summary.get("optimizer_step"))
    return trainer


def record_train_gate(ledger, trainer, summary) -> dict:
    """Record the 'train' gate for what `trainer.train()` ACTUALLY did.

    THE STOPPED-RECORDED-AS-PASSED DEFECT (H2 review, 2026-10-13). `train()`
    returns normally after a COORDINATED STOP (SIGTERM, or train.max_wall_seconds)
    -- it checkpoints and returns -- and the gate used to be recorded
    `passed=True` unconditionally. `main` then (correctly) skipped the eval and
    exited 128+n / 82, but the ledger on disk now certified a half-trained
    checkpoint as a finished 'train' stage: anything that reads the ledger
    instead of the exit status (a resumed `evaluate_and_gate`, an operator
    reading `stage_gates.json`) was told the stage had trained.

    A stopped run is neither finished nor failed: it is RESUMABLE. The ledger
    has no third state, so it is recorded `passed=False` with a `STOPPED
    (resumable)` detail; a rerun resumes from the checkpoint and records the
    real verdict at the end of the real run.
    """
    step = summary.get("optimizer_step")
    if getattr(trainer, "stop_exit_code", 0):
        return ledger.record(
            "train",
            passed=False,
            detail=f"STOPPED (resumable): {trainer.stop_reason} stop at optimizer_step={step}; "
            "checkpointed, not finished",
        )
    return ledger.record("train", passed=True, detail=str(step))


def main(argv=None) -> int:
    args = parse_args(argv)

    # FIRST: every env var torch reads at import, plus the heavy-job flock and
    # the /apps path assertions. `heavy=True` also runs the host preflight;
    # `rss_cfg` adds the full-FT RSS projection on top of the generic floor.
    cfg = load_config(args.config, args.overrides)
    apply_resource_guards(cpu_threads=8, heavy=True, label="sft-stage-11", rss_cfg=cfg)

    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
    )

    # BEFORE anything touches CUDA. THE DUPLICATE-GPU DEFECT: torchrun starts
    # four processes but assigns no devices, so without this every rank's
    # `.cuda()` lands on device 0 and NCCL dies with "Duplicate GPU detected:
    # rank 2 and rank 0 both on CUDA device 1b0" -> ncclInvalidUsage. Section-01
    # wrote `init_distributed` for exactly this (init_process_group +
    # set_device(local_rank)); this script simply never called it.
    from align.utils.dist import init_distributed

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    # run_dir: the NCCL flight-recorder dump lands under it (never /tmp).
    init_distributed(run_dir=run_dir)

    # Microseconds, and it converts the whole duplicate-GPU class into a named
    # launch failure instead of an ncclInvalidUsage 400 lines into a forward.
    from align.train.stage_gates import assert_distinct_devices

    assert_distinct_devices()
    ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())

    banner = run_gates(cfg, ledger)
    from align.registry import render_banner

    print(render_banner(banner, title=f"align SFT stage :: {cfg.run_name}"))

    model = build_model(cfg)
    trainer = train(cfg, ledger, model, banner)
    if trainer.stop_exit_code:
        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved and
        # returned: the stage is NOT finished, so do not fall into its
        # post-train eval or announce completion. 82 = wall budget (the
        # supervisor will not relaunch it), 128+n = the signal.
        print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr)
        return trainer.stop_exit_code

    # BARRIER, then RELEASE, then eval -- both are load-bearing:
    # 1. train() returns minutes earlier on ranks 1..N-1 than on rank 0 (the
    #    sole checkpoint writer). Without this barrier they raced into the
    #    eval, OOM'd building a second 7B, and torchrun SIGTERM'd rank 0
    #    MID-SAVE -- whose cleanup handler then deleted the .partial. The
    #    epoch's tail died twice this way.
    # 2. run_stage_eval builds its OWN model from the checkpoint; the
    #    training replica + optimizer must be freed first or the engine's
    #    KV-cache preflight refuses (1.2GB free vs 4.8GB required).
    from align.utils.dist import barrier

    barrier()
    checkpoint = trainer.latest_checkpoint()
    if cfg.model.arch != "mistral":
        # The post-train MATH-500/GSM8K gate certifies the Mistral stack; a qwen3 checkpoint is
        # judged by the Lean loop. Say so instead of silently skipping.
        print(f"SFT stage complete (arch={cfg.model.arch}; no MATH/GSM8K gate applies). Checkpoint: {checkpoint}")
        return 0
    trainer.release()  # model, optimizer AND the probe's references (U01-11)
    del model
    import gc

    import torch

    gc.collect()
    torch.cuda.empty_cache()
    decision = evaluate_and_gate(cfg, ledger, checkpoint=checkpoint)
    if not decision.passed:
        print(decision.halt_message, file=sys.stderr)
        return 1

    print(f"SFT stage complete. DPO base checkpoint: {checkpoint}")
    return 0


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

    raise SystemExit(run_with_exit_codes(main))