"""``softasp-train`` — launch a training run (supports torchrun for DDP).

Usage::

    # Single-GPU
    uv run --locked softasp-train instantiation=sudoku model.pretrained_from=backbone.pt

    # Multi-GPU via torchrun
    uv run --locked torchrun --nproc_per_node=4 -m softasp.cli.train \\
        instantiation=my_distributed_task paths=local
"""

from __future__ import annotations

import logging
from pathlib import Path

import hydra
from omegaconf import DictConfig, OmegaConf
from softasp_core.training.schedule import CORRECTED_PROTOCOL, LEGACY_PROTOCOL
from softasp_core.training.seed import seed_everything
from softasp_core.training.trainer import Trainer
from softasp_core.utils.logging import setup_logging

from softasp.configuration import config_directory
from softasp.instantiations.checkpoint import read_checkpoint, validate_checkpoint
from softasp.instantiations.runtime import provider

logger = logging.getLogger(__name__)

_CONFIGS = config_directory()


@hydra.main(version_base="1.3", config_path=_CONFIGS, config_name="config")
def main(cfg: DictConfig) -> None:
    run_training(cfg)


def run_training(cfg: DictConfig) -> Path:
    """Train from an explicit configuration and return the exact run directory."""
    setup_logging(level=cfg.logging.get("level", "INFO"))
    logger.info("=== softasp-train ===")

    # Resolve all interpolations
    OmegaConf.resolve(cfg)

    protocol = cfg.trainer.get("protocol", LEGACY_PROTOCOL)
    if protocol not in (LEGACY_PROTOCOL, CORRECTED_PROTOCOL):
        raise ValueError(f"Unknown training protocol: {protocol}")
    if protocol == CORRECTED_PROTOCOL:
        # Includes release-model construction and all subsequent stochastic setup.
        seed_everything(cfg.trainer.seed, deterministic=cfg.trainer.deterministic)
        OmegaConf.update(cfg, "initialization_seeded_before_model", True, force_add=True)

    model_provider = provider(cfg, "model")
    # Preparation precedes data construction to preserve initialization order.
    prepared = model_provider.prepare(cfg)
    train_ds, val_ds = provider(cfg, "dataset").build(cfg)
    model = prepared if prepared is not None else model_provider.build(cfg)
    evaluator = provider(cfg, "metrics").build(cfg)

    # Train
    trainer = Trainer(
        cfg=cfg, model=model, evaluator=evaluator,
        checkpoint_validator=lambda path: validate_checkpoint(cfg, read_checkpoint(path)),
    )
    run_dir = trainer.fit(train_ds, val_ds)

    logger.info("Training complete.")
    return run_dir


if __name__ == "__main__":
    main()