#!/usr/bin/env python
"""Pretraining stage entry (section-14): gates -> wrap -> train -> checkpoint.

    torchrun --nproc_per_node=4 scripts/pretrain.py configs/pretrain.toml --optimizer.lr=3e-5

The Pretraining pillar, OLMoE-shaped and honestly toy-scaled: continued
pretraining of the base model on a capped (~50M token), STREAMED slice of a
raw-text corpus, packed into the repo's shard format by
align/data/pretrain_corpus.py. The objective is the SAME packed next-token NLL
the SFT stage computes -- every pretrain row stores prompt_len 0, so
`collate_sft` labels every token (labels == input_ids, no masking) and
`make_sft_loss_adapter` runs unchanged. One loss path, two stages.

Thin by design, mirroring scripts/sft.py: everything this drives is built and
tested elsewhere; what lives here is ORDER:

    build_currency -> smoke_7b -> probe_20step -> rss_preflight -> train

Differences from the SFT entry, each deliberate:
  * NO golden-regression bridge. The golden anchor certifies the SFT eval
    stack (MATH-500/GSM8K extraction against known scores); a base-style
    pretrain checkpoint has no such anchor. configs/pretrain.toml sets
    gates.require_golden_receipt: false, which run_entry_gates honors with a
    loud warning rather than silence.
  * NO accuracy quality gate at exit. The stage's product is a checkpoint and
    its train-NLL trajectory; downstream SFT is its consumer and its judge.

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.
tests/test_pretrain_entry.py enforces this by AST, exactly as the SFT test
does for scripts/sft.py.

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 `python scripts/pretrain.py` (the documented
# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py).
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.config import load_config
from align.train.stage_gates import (
    GateLedger,
    assert_build_is_current,
    assert_probe_resolved,
    eval_source_hash,
    parse_args,
    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."""
    from align.train.stage_smoke import build_stage_model

    return build_stage_model(cfg)


def run_gates(cfg, ledger):
    """The entry sequence. Refuses to proceed on a stale or unstamped corpus
    build before it touches a GPU. The golden bridge is intentionally absent
    (see module docstring); run_entry_gates still calls require_golden_receipt
    internally, which configs/pretrain.toml downgrades to a warning via
    gates.require_golden_receipt: false."""
    assert_build_is_current(cfg.data.manifest_path, cfg)
    banner = run_entry_gates(cfg, ledger)
    assert_probe_resolved(banner)
    return banner


def train(cfg, ledger, model, banner):
    """One pretraining epoch through the section-07 Trainer.

    Identical wiring to the SFT stage on purpose: same sampler construction
    (the SAME sampler the entry gates measured the horizon from), same
    collate, same fused loss adapter. The pretraining objective difference
    lives entirely in the data -- prompt_len 0 makes collate_sft emit
    labels == input_ids."""
    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

    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))

    trainer = Trainer(
        cfg=cfg,
        model=model,
        sampler=sampler,
        loss_adapter=adapter,
        run_dir=os.path.join(cfg.paths.checkpoint_dir, "pretrain"),
        stage="pretrain",
    )
    trainer.resume_from_latest()
    summary = trainer.train()
    record_train_gate(ledger, trainer, summary)
    return trainer


def record_train_gate(ledger, trainer, summary) -> dict:
    """Record the 'train' gate for what `trainer.train()` ACTUALLY did (same contract as
    scripts/sft.py::record_train_gate; duplicated because scripts/ is not a package).

    THE STOPPED-RECORDED-AS-PASSED DEFECT (H2 review, 2026-10-13; the sibling the first
    fix missed). `train()` returns normally after a COORDINATED STOP (SIGTERM or
    train.max_wall_seconds) and this used to record `passed=True` unconditionally, so the
    ledger certified a half-trained checkpoint as a finished stage while `main` exited
    128+n / 82. A stopped run is resumable, neither finished nor failed; the ledger has no
    third state, so it is `passed=False` with a `STOPPED (resumable)` detail, and the real
    verdict overwrites it when the resumed run finishes.
    """
    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="pretrain-stage-14", 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 (see
    # scripts/sft.py): torchrun assigns no devices, so every rank's `.cuda()`
    # lands on device 0 unless init_distributed sets one per local rank.
    from align.utils.dist import init_distributed

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

    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 pretrain 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

    checkpoint = trainer.latest_checkpoint()
    print(f"Pretrain stage complete. SFT base checkpoint: {checkpoint}")
    return 0


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

    raise SystemExit(run_with_exit_codes(main))