from __future__ import annotations

import pytest

from argus.ray.builders import build_training_argv
from argus.ray.models import JobSpec


def _base_spec(workload: str, *, total_training_steps: int | None = 10) -> JobSpec:
    return JobSpec(
        workload=workload,
        submission_id=None,
        code_path="/code",
        model_id="m",
        train_file="train.jsonl",
        val_file="val.jsonl",
        nnodes=2,
        n_gpus_per_node=4,
        total_epochs=1,
        total_training_steps=total_training_steps,
        save_freq=10,
        test_freq=None,
        trainer_device=None,
    )


def test_build_training_argv_ppo_smoke():
    spec = _base_spec("ppo")
    built = build_training_argv(spec, submission_id="sid", job_dir="/job")
    assert built.argv[:3] == ["python3", "-m", "verl.trainer.main_ppo"]
    assert "data.val_files=val.jsonl" in built.argv
    # v3.7: default rollout backend switches from sglang -> vllm.
    assert "actor_rollout_ref.rollout.name=vllm" in built.argv
    assert "actor_rollout_ref.rollout.name=sglang" not in built.argv
    assert "trainer.test_freq=-1" in built.argv


def test_build_training_argv_grpo_has_override():
    spec = _base_spec("grpo")
    built = build_training_argv(spec, submission_id="sid", job_dir="/job")
    assert "algorithm.adv_estimator=grpo" in built.argv


def test_build_training_argv_sft_smoke():
    spec = _base_spec("sft")
    built = build_training_argv(spec, submission_id="sid", job_dir="/job")
    assert built.argv[:3] == ["python3", "-m", "verl.trainer.sft_trainer_ray"]
    assert "trainer.device=cpu" in built.argv
    assert "data.val_files=val.jsonl" in built.argv


def test_build_training_argv_omits_total_training_steps_when_unset():
    spec = _base_spec("ppo", total_training_steps=None)
    built = build_training_argv(spec, submission_id="sid", job_dir="/job")
    assert not any(a.startswith("trainer.total_training_steps=") for a in built.argv)


def test_build_training_argv_unsupported_raises():
    spec = _base_spec("bad")
    with pytest.raises(ValueError, match="unsupported workload"):
        build_training_argv(spec, submission_id="sid", job_dir="/job")


def test_build_advanced_argv_uses_bash_lc():
    from argus.ray.builders import build_advanced_argv

    built = build_advanced_argv("echo hi")
    assert built.argv[:2] == ["bash", "-lc"]
