from __future__ import annotations

import pytest

from argus.service.advanced_command import expand_advanced_command, validate_advanced_command


def test_expand_maps_home_common_datasets_to_private_datasets():
    cmd = "python3 -m verl.trainer.main_ppo data.train_files=$HOME/common/datasets/gsm8k/train.parquet"
    expanded = expand_advanced_command(user_id="alice", command=cmd)
    assert "/private/datasets/gsm8k/train.parquet" in expanded
    assert "$HOME/common/datasets" not in expanded


def test_expand_maps_home_common_hf_to_private_hf():
    cmd = "python3 -m verl.trainer.main_ppo HF_HOME=$HOME/common/hf"
    expanded = expand_advanced_command(user_id="alice", command=cmd)
    assert "HF_HOME=/private/hf" in expanded


def test_expand_maps_home_to_private_users():
    cmd = "python3 -m verl.trainer.main_ppo data.train_files=$HOME/datasets/x.parquet"
    expanded = expand_advanced_command(user_id="alice", command=cmd)
    assert "data.train_files=/private/users/alice/datasets/x.parquet" in expanded


def test_validate_requires_verl_trainer_entry():
    with pytest.raises(ValueError, match="-m verl\\."):
        validate_advanced_command(user_id="alice", expanded_command="echo hi")


def test_validate_rejects_cross_user_paths():
    cmd = "python3 -m verl.trainer.main_ppo x=/private/users/bob/datasets/x"
    with pytest.raises(ValueError, match="another user's path"):
        validate_advanced_command(user_id="alice", expanded_command=cmd)


def test_validate_allows_train_val_under_allowed_roots():
    cmd = (
        "python3 -m verl.trainer.main_ppo "
        "data.train_files=/private/datasets/gsm8k/train.parquet "
        "data.val_files=/private/users/alice/datasets/gsm8k/test.parquet"
    )
    validate_advanced_command(user_id="alice", expanded_command=cmd)


def test_validate_rejects_train_val_outside_allowed_roots():
    cmd = (
        "python3 -m verl.trainer.main_ppo "
        "data.train_files=/etc/passwd "
        "data.val_files=/private/datasets/gsm8k/test.parquet"
    )
    with pytest.raises(ValueError, match="train_file must be under allowed roots"):
        validate_advanced_command(user_id="alice", expanded_command=cmd)


def test_validate_allows_reward_path_under_user_code():
    cmd = (
        "python3 -m verl.trainer.main_ppo "
        "custom_reward_function.path=/private/users/alice/code/reward.py "
        "custom_reward_function.name=compute_score"
    )
    validate_advanced_command(user_id="alice", expanded_command=cmd)


def test_validate_rejects_reward_path_outside_user_code():
    cmd = (
        "python3 -m verl.trainer.main_ppo "
        "custom_reward_function.path=/private/datasets/reward.py "
        "custom_reward_function.name=compute_score"
    )
    with pytest.raises(ValueError, match="reward path must be under"):
        validate_advanced_command(user_id="alice", expanded_command=cmd)
