#!/usr/bin/env python
from __future__ import annotations
import argparse
from pathlib import Path
import yaml
from sacflow.utils.config import load_yaml, save_yaml, set_by_path
from sacflow.utils.misc import seed_everything, ensure_dir
from sacflow.utils.distributed import init_distributed, cleanup, is_main_process, barrier
from sacflow.utils.wandb_utils import init_wandb, wandb_finish
from sacflow.engine.train_loop import run_training


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--name", default=None)
    ap.add_argument("--resume", nargs="?", const="auto", default=None, help="Resume checkpoint path, or use --resume without a value for auto-detection.")
    ap.add_argument("--opts", nargs="*", default=[], help="Override config values: key=value, e.g. train.epochs=220 optim.lr=1e-4")
    args = ap.parse_args()
    cfg = load_yaml(args.config)
    for opt in args.opts:
        if "=" not in opt:
            raise ValueError(f"Invalid --opts entry {opt!r}; expected key=value")
        key, value = opt.split("=", 1)
        try:
            parsed = yaml.safe_load(value)
        except Exception:
            parsed = value
        set_by_path(cfg, key, parsed)
    if args.resume is not None:
        cfg.setdefault("train", {})["resume_checkpoint"] = args.resume
    seed_everything(int(cfg.get("seed", 1337)))
    device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
    out_dir = ensure_dir(cfg["output_dir"])
    if is_main_process():
        save_yaml(cfg, out_dir / "resolved_config.yaml")
    run = init_wandb(cfg, run_name=args.name or Path(cfg["output_dir"]).name)
    try:
        run_training(cfg, device, wandb_run=run)
    finally:
        barrier()
        wandb_finish(run)
        cleanup()

if __name__ == "__main__":
    main()