# SPDX-License-Identifier: Apache-2.0
"""FlashGen YAML training entrypoint (NPU-first).

Thin wrapper over FastVideo's training loop. Its only job is to pin an
NPU-safe attention backend (Torch SDPA) before any FastVideo module reads it,
then delegate to :func:`fastvideo.train.entrypoint.train.run_training_from_config`.

Usage::

    torchrun --nproc_per_node=<N> -m flashgen.entrypoint \
        --config flashgen/configs/wan_dmd_npu.yaml
"""

from __future__ import annotations

import argparse
import os
import sys
from typing import Any


def _patch_cuda_rng_for_npu() -> None:
    """Proxy ``torch.cuda`` RNG helpers to NPU when CUDA is unavailable.

    FastVideo's checkpoint code calls ``torch.cuda.get_rng_state()`` /
    ``torch.cuda.set_rng_state()`` unconditionally when saving/restoring a
    training-state checkpoint. On an Ascend build (torch without CUDA) these
    raise ``AssertionError: Torch not compiled with CUDA enabled`` and crash
    every rank at the first checkpoint step. Redirect them to the ``torch.npu``
    equivalents (falling back to a no-op) so checkpointing works on NPU.
    """
    import torch

    if torch.cuda.is_available():
        return

    npu_mod = getattr(torch, "npu", None)
    npu_get = getattr(npu_mod, "get_rng_state", None)
    npu_set = getattr(npu_mod, "set_rng_state", None)

    def _get_rng_state(*_args: Any, **_kwargs: Any) -> "torch.Tensor":
        if callable(npu_get):
            try:
                return npu_get()
            except Exception:
                # NPU RNG unavailable; fall back to empty placeholder below.
                return torch.empty(0, dtype=torch.uint8)
        # Placeholder so the snapshot stays self-consistent (set skips it).
        return torch.empty(0, dtype=torch.uint8)

    def _set_rng_state(new_state: Any, *_args: Any, **_kwargs: Any) -> None:
        if not callable(npu_set):
            return
        # Skip the empty placeholder produced when NPU RNG is unavailable.
        if isinstance(new_state, torch.Tensor) and new_state.numel() == 0:
            return
        try:
            npu_set(new_state)
        except Exception:
            return

    torch.cuda.get_rng_state = _get_rng_state
    torch.cuda.set_rng_state = _set_rng_state


def _setup_npu_environment() -> None:
    """Pin an NPU-safe attention backend and enable Ascend NPU if present.

    ``os.environ.setdefault`` is used so an explicit user override (e.g. via a
    ``--config`` override or a pre-set env var) always wins.
    """
    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")

    # Best-effort Ascend NPU activation. Importing torch_npu registers the
    # ``npu`` device backend with PyTorch. It is intentionally optional so the
    # package still imports on CPU/CUDA machines.
    try:
        import torch_npu  # type: ignore  # noqa: F401
    except Exception:
        return

    # torch_npu is present: make FastVideo's CUDA-only RNG checkpoint calls
    # work on Ascend by proxying them to the NPU RNG helpers.
    _patch_cuda_rng_for_npu()


def run_training_from_config(
    config_path: str,
    *,
    dry_run: bool = False,
    overrides: list[str] | None = None,
) -> None:
    """Set up the NPU environment, then run FastVideo's training loop."""
    _setup_npu_environment()

    from fastvideo.train.entrypoint.train import (
        run_training_from_config as _fv_run_training_from_config,
    )

    _fv_run_training_from_config(
        config_path,
        dry_run=dry_run,
        overrides=overrides,
    )


def main(
    args: Any,
    overrides: list[str] | None = None,
) -> None:
    run_training_from_config(
        str(args.config),
        dry_run=bool(args.dry_run),
        overrides=overrides,
    )


def cli(argv: list[str] | None = None) -> None:
    """Console-script entry point (parses arguments from ``sys.argv``)."""
    parser = argparse.ArgumentParser(
        description="FlashGen YAML training entrypoint (NPU-first).",
    )
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to training YAML config.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help=("Parse config and build runtime, but do not start training."),
    )
    args, unknown = parser.parse_known_args(sys.argv[1:] if argv is None else argv)
    main(args, overrides=unknown if unknown else None)


if __name__ == "__main__":
    cli()