from __future__ import annotations

import re


_RE_PRIVATE_USER = re.compile(r"/private/users/([^/\s\"']+)")
_RE_TRAIN_FILES = re.compile(r"data\.train_files=([^\s\\]+)")
_RE_VAL_FILES = re.compile(r"data\.val_files=([^\s\\]+)")
_RE_REWARD_PATH = re.compile(r"custom_reward_function\.path=([^\s\\]+)")


def _strip_quotes(s: str) -> str:
    s = s.strip()
    if len(s) >= 2 and s[0] == s[-1] and s[0] in ("'", '"'):
        return s[1:-1]
    return s


def expand_advanced_command(*, user_id: str, command: str) -> str:
    """
    Expand user-facing macros into real container paths.

    Important: `$HOME/common/...` is a SFTPGo virtual folder; the training process
    cannot see it. We map those paths to real shared mounts.
    """

    home = f"/private/users/{user_id}"
    expanded = command
    expanded = expanded.replace("$HOME/common/datasets", "/private/datasets")
    expanded = expanded.replace("$HOME/common/hf", "/private/hf")
    expanded = expanded.replace("$HOME", home)
    return expanded


def validate_advanced_command(*, user_id: str, expanded_command: str) -> None:
    """
    Best-effort validation for advanced commands.

    Strong constraints (raise ValueError on failure):
    - Must include a python3 VERL trainer entrypoint
    - Must not reference other users' /private/users/<other>/... paths
    - If train/val files are present, must be under allowed roots
    - If reward path is present, must be under /private/users/<me>/code/...
    """

    # v3.5: advanced commands must still be scoped to VERL entrypoints (not arbitrary shell).
    # Allow non-trainer utilities like `verl.model_merger`.
    if "python3" not in expanded_command or "-m verl." not in expanded_command:
        raise ValueError("command must include 'python3' and '-m verl.'")

    for other in _RE_PRIVATE_USER.findall(expanded_command):
        if other != user_id:
            raise ValueError(f"command references another user's path: /private/users/{other}/...")

    allowed_dataset_prefixes = (
        f"/private/users/{user_id}/datasets/",
        "/private/datasets/",
        # Backward compatible with v3.0 paths used by existing templates/users.
        "/private/common/datasets/",
    )

    for m in _RE_TRAIN_FILES.finditer(expanded_command):
        path = _strip_quotes(m.group(1))
        if not path.startswith(allowed_dataset_prefixes):
            raise ValueError(f"train_file must be under allowed roots, got: {path}")

    for m in _RE_VAL_FILES.finditer(expanded_command):
        path = _strip_quotes(m.group(1))
        if not path.startswith(allowed_dataset_prefixes):
            raise ValueError(f"val_file must be under allowed roots, got: {path}")

    for m in _RE_REWARD_PATH.finditer(expanded_command):
        path = _strip_quotes(m.group(1))
        allow = f"/private/users/{user_id}/code/"
        if not path.startswith(allow):
            raise ValueError(f"reward path must be under {allow}, got: {path}")
