#!/usr/bin/env python3
"""Run causal LM pretraining (config-driven); optional resume from a checkpoint."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from weather_llm.config_loader import repo_root_from
from weather_llm.training.pretrain import run_pretrain


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--pretrain_config", type=Path, default=None)
    ap.add_argument("--model_config", type=Path, default=None)
    ap.add_argument(
        "--resume_from_checkpoint",
        type=Path,
        default=None,
        help="Trainer checkpoint-* dir under output_dir (has trainer_state.json); overrides YAML",
    )
    args = ap.parse_args()
    root = repo_root_from(ROOT)
    pc = args.pretrain_config or (root / "configs/pretrain_prototype.yaml")
    mc = args.model_config or (root / "configs/model_prototype.yaml")
    run_pretrain(
        pc.resolve(),
        mc.resolve(),
        resume_from_checkpoint=args.resume_from_checkpoint,
    )
    return 0


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